From b46b96d7dd7be8f3b7819649977da42fe23931a0 Mon Sep 17 00:00:00 2001 From: vasilyl Date: Tue, 2 Jun 2026 03:36:02 +0000 Subject: [PATCH 1/3] feat: Add daily time partitioning to BigQuery event logging table --- .../bigquery_agent_analytics_plugin.go | 4 + .../bigquery_agent_analytics_plugin_test.go | 85 +++++++++++++++++++ 2 files changed, 89 insertions(+) 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) + } +} From 3f7fdbe5ab726c508eebe59cd0dab86f317e92c6 Mon Sep 17 00:00:00 2001 From: vasilyl Date: Tue, 2 Jun 2026 04:01:00 +0000 Subject: [PATCH 2/3] feat: Add ConnectionID for BigQuery Object Tables secure external access configuration --- plugin/agentanalytics/config.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/plugin/agentanalytics/config.go b/plugin/agentanalytics/config.go index 40e3b5595..42c05f0d7 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: "", } } From 88e9cf81a24e3c35d03792a398e77ee958745398 Mon Sep 17 00:00:00 2001 From: vasilyl Date: Tue, 2 Jun 2026 20:12:42 +0000 Subject: [PATCH 3/3] style: Fix gofumpt linter issues in object-tables-connection branch --- plugin/agentanalytics/config.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugin/agentanalytics/config.go b/plugin/agentanalytics/config.go index 42c05f0d7..871e76849 100644 --- a/plugin/agentanalytics/config.go +++ b/plugin/agentanalytics/config.go @@ -89,6 +89,6 @@ func DefaultConfig() Config { MaxDelay: 10 * time.Second, Multiplier: 2.0, }, - ConnectionID: "", + ConnectionID: "", } }