diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go index 4b5822e17..c03dadd2f 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go +++ b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go @@ -102,6 +102,10 @@ func NewBigQueryAgentAnalyticsPluginWithClients( } err = tableRef.Create(ctx, &bq.TableMetadata{ Schema: EventsSchema(), + TimePartitioning: &bq.TimePartitioning{ + Field: "timestamp", + Type: bq.DayPartitioningType, + }, Clustering: &bq.Clustering{ Fields: config.ClusteringFields, }, diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go index 2a59eefcf..316f45a82 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go +++ b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go @@ -15,6 +15,7 @@ package agentanalytics import ( + "bytes" "context" "errors" "io" @@ -342,3 +343,87 @@ func TestLogEvent_ExtractsTraceInfo(t *testing.T) { t.Error("Timed out waiting for request") } } + +func TestNewBigQueryAgentAnalyticsPlugin_CreateTable_WithPartitioning(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + + createCalled := false + var requestBody string + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + // Table metadata request: returns 404 Not Found to trigger creation + if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") { + return &http.Response{ + StatusCode: http.StatusNotFound, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":404,"message":"Not found"}}`)), + }, nil + } + // Table creation request + if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") { + createCalled = true + bodyBytes, _ := io.ReadAll(r.Body) + requestBody = string(bodyBytes) + r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("{}")), + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("{}")), + }, nil + }, + } + httpClient := &http.Client{Transport: mockTransport} + bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient)) + if err != nil { + t.Fatalf("Failed to create bigquery client: %v", err) + } + + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + gSrv := grpc.NewServer() + storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{}) + go func() { _ = gSrv.Serve(lis) }() + t.Cleanup(gSrv.Stop) + + conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to dial test server: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn)) + if err != nil { + t.Fatalf("Failed to create BigQuery write client: %v", err) + } + + _, err = NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err != nil { + t.Fatalf("Plugin initialization error: %v", err) + } + + if !createCalled { + t.Error("Expected table creation to be called") + } + + if !strings.Contains(requestBody, "timePartitioning") { + t.Errorf("Expected request body to contain 'timePartitioning', got: %s", requestBody) + } + if !strings.Contains(requestBody, "DAY") { + t.Errorf("Expected partitioning type to be 'DAY', got request body: %s", requestBody) + } + if !strings.Contains(requestBody, "timestamp") { + t.Errorf("Expected partitioning field to be 'timestamp', got request body: %s", requestBody) + } +} diff --git a/plugin/agentanalytics/config.go b/plugin/agentanalytics/config.go index 40e3b5595..871e76849 100644 --- a/plugin/agentanalytics/config.go +++ b/plugin/agentanalytics/config.go @@ -63,6 +63,9 @@ type Config struct { // Retry configuration for appending rows. RetryConfig RetryConfig + + // BigQuery connection ID for ObjectRef secure external access (location.connection_id) + ConnectionID string } // DefaultConfig returns the default configuration for the agent analytics plugin. @@ -86,5 +89,6 @@ func DefaultConfig() Config { MaxDelay: 10 * time.Second, Multiplier: 2.0, }, + ConnectionID: "", } }