diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go index 4b5822e17..652e3abb5 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go +++ b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go @@ -18,11 +18,14 @@ package agentanalytics import ( "context" "fmt" + "net/http" + "strings" "time" bq "cloud.google.com/go/bigquery" bqstorage "cloud.google.com/go/bigquery/storage/apiv1" "go.opentelemetry.io/otel/trace" + "google.golang.org/api/googleapi" "google.golang.org/genai" "google.golang.org/adk/v2/agent" @@ -32,6 +35,98 @@ import ( "google.golang.org/adk/v2/tool" ) +const ( + schemaVersion = "1" + schemaVersionLabelKey = "adk_schema_version" + viewSQLTemplate = "CREATE OR REPLACE VIEW `%s.%s.%s` AS\nSELECT\n %s\nFROM\n `%s.%s.%s`\nWHERE\n event_type = '%s'" +) + +var viewCommonColumns = []string{ + "timestamp", + "event_type", + "agent", + "session_id", + "invocation_id", + "user_id", + "trace_id", + "span_id", + "parent_span_id", + "status", + "error_message", + "is_truncated", +} + +var eventViewDefs = map[string][]string{ + "USER_MESSAGE": {}, + "MODEL_REQUEST": { + "JSON_VALUE(attributes, '$.model') AS model", + "content AS request_content", + "JSON_QUERY(attributes, '$.llm_config') AS llm_config", + "JSON_QUERY(attributes, '$.tools') AS tools", + }, + "MODEL_RESPONSE": { + "JSON_QUERY(content, '$.response') AS response", + "CAST(JSON_VALUE(content, '$.usage.prompt') AS INT64) AS usage_prompt_tokens", + "CAST(JSON_VALUE(content, '$.usage.completion') AS INT64) AS usage_completion_tokens", + "CAST(JSON_VALUE(content, '$.usage.total') AS INT64) AS usage_total_tokens", + "CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms", + "CAST(JSON_VALUE(latency_ms, '$.time_to_first_token_ms') AS INT64) AS ttft_ms", + "JSON_VALUE(attributes, '$.model_version') AS model_version", + "JSON_QUERY(attributes, '$.usage_metadata') AS usage_metadata", + }, + "MODEL_ERROR": { + "CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms", + }, + "TOOL_START": { + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.args') AS tool_args", + "JSON_VALUE(content, '$.tool_origin') AS tool_origin", + }, + "TOOL_END": { + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.result') AS tool_result", + "JSON_VALUE(content, '$.tool_origin') AS tool_origin", + "CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms", + }, + "TOOL_ERROR": { + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.args') AS tool_args", + "JSON_VALUE(content, '$.tool_origin') AS tool_origin", + "CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms", + }, + "AGENT_START": { + "JSON_VALUE(content, '$.text_summary') AS agent_instruction", + }, + "AGENT_END": { + "CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms", + }, + "INVOCATION_START": {}, + "INVOCATION_END": {}, + "EVENT": {}, + "STATE_DELTA": { + "JSON_QUERY(attributes, '$.state_delta') AS state_delta", + }, + "HITL_CREDENTIAL_REQUEST": { + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.args') AS tool_args", + }, + "HITL_CONFIRMATION_REQUEST": { + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.args') AS tool_args", + }, + "HITL_INPUT_REQUEST": { + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.args') AS tool_args", + }, + "A2A_INTERACTION": { + "content AS response_content", + "JSON_VALUE(attributes, '$.a2a_metadata.\"a2a:task_id\"') AS a2a_task_id", + "JSON_VALUE(attributes, '$.a2a_metadata.\"a2a:context_id\"') AS a2a_context_id", + "JSON_QUERY(attributes, '$.a2a_metadata.\"a2a:request\"') AS a2a_request", + "JSON_QUERY(attributes, '$.a2a_metadata.\"a2a:response\"') AS a2a_response", + }, +} + // NewBigQueryAgentAnalyticsPlugin creates a newly configured analytics plugin with default config. func NewBigQueryAgentAnalyticsPlugin( ctx context.Context, @@ -95,21 +190,42 @@ func NewBigQueryAgentAnalyticsPluginWithClients( // Ensure table exists (done once per instance) tableRef := bqClient.Dataset(config.DatasetID).Table(config.TableName) - _, err := tableRef.Metadata(ctx) + meta, err := tableRef.Metadata(ctx) if err != nil { - if config.Logger != nil { - config.Logger.Printf("Table %s not found. Creating it...", config.TableName) - } - err = tableRef.Create(ctx, &bq.TableMetadata{ - Schema: EventsSchema(), - Clustering: &bq.Clustering{ - Fields: config.ClusteringFields, - }, - }) - if err != nil { + if apiErr, ok := err.(*googleapi.Error); ok && apiErr.Code == http.StatusNotFound { if config.Logger != nil { - config.Logger.Printf("Failed to create BigQuery table %v: %v", config.TableName, err) + config.Logger.Printf("Table %s not found. Creating it...", config.TableName) + } + err = tableRef.Create(ctx, &bq.TableMetadata{ + Schema: EventsSchema(), + TimePartitioning: &bq.TimePartitioning{ + Field: "timestamp", + Type: bq.DayPartitioningType, + }, + Clustering: &bq.Clustering{ + Fields: config.ClusteringFields, + }, + Labels: map[string]string{ + schemaVersionLabelKey: schemaVersion, + }, + }) + if err != nil { + return nil, fmt.Errorf("failed to create BigQuery table: %w", err) } + } else { + return nil, fmt.Errorf("failed to retrieve BigQuery table metadata: %w", err) + } + } else { + if config.AutoSchemaUpgrade { + if err := maybeUpgradeSchema(ctx, tableRef, meta, config.Logger); err != nil { + return nil, fmt.Errorf("failed to auto-upgrade BigQuery table schema: %w", err) + } + } + } + + if config.CreateViews { + if err := createAnalyticsViews(ctx, bqClient, config); err != nil && config.Logger != nil { + config.Logger.Printf("Views creation failed: %v", err) } } @@ -297,3 +413,156 @@ func NewBigQueryAgentAnalyticsPluginWithClients( return baseplugin.New(cfg) } + +func schemaFieldsMatch(existing, desired bq.Schema) (newFields, updatedRecords bq.Schema) { + existingByName := make(map[string]*bq.FieldSchema) + for _, f := range existing { + existingByName[f.Name] = f + } + + for _, desiredField := range desired { + existingField, exists := existingByName[desiredField.Name] + if !exists { + newFields = append(newFields, desiredField) + } else if desiredField.Type == bq.RecordFieldType && existingField.Type == bq.RecordFieldType && len(desiredField.Schema) > 0 { + subNew, subUpdated := schemaFieldsMatch(existingField.Schema, desiredField.Schema) + if len(subNew) > 0 || len(subUpdated) > 0 { + // Build a merged sub-field list + mergedSub := make(bq.Schema, len(existingField.Schema)) + copy(mergedSub, existingField.Schema) + + updatedNames := make(map[string]bool) + for _, f := range subUpdated { + updatedNames[f.Name] = true + } + + // Replace updated nested records in-place + for i, f := range mergedSub { + if updatedNames[f.Name] { + for _, u := range subUpdated { + if u.Name == f.Name { + mergedSub[i] = u + break + } + } + } + } + + // Append entirely new sub-fields + mergedSub = append(mergedSub, subNew...) + + // Create updated RECORD field schema + updatedField := *existingField + updatedField.Schema = mergedSub + updatedRecords = append(updatedRecords, &updatedField) + } + } + } + return newFields, updatedRecords +} + +func maybeUpgradeSchema(ctx context.Context, existingTable *bq.Table, meta *bq.TableMetadata, logger Logger) error { + storedVersion := meta.Labels[schemaVersionLabelKey] + if storedVersion == schemaVersion { + return nil + } + + newFields, updatedRecords := schemaFieldsMatch(meta.Schema, EventsSchema()) + + if len(newFields) > 0 || len(updatedRecords) > 0 { + updatedNames := make(map[string]bool) + for _, f := range updatedRecords { + updatedNames[f.Name] = true + } + + merged := make(bq.Schema, 0, len(meta.Schema)+len(newFields)) + for _, f := range meta.Schema { + if updatedNames[f.Name] { + for _, u := range updatedRecords { + if u.Name == f.Name { + merged = append(merged, u) + break + } + } + } else { + merged = append(merged, f) + } + } + merged = append(merged, newFields...) + + // Update schema in metadata update + update := bq.TableMetadataToUpdate{ + Schema: merged, + } + update.SetLabel(schemaVersionLabelKey, schemaVersion) + + _, err := existingTable.Update(ctx, update, meta.ETag) + if err != nil { + return fmt.Errorf("failed to update table schema: %w", err) + } + + if logger != nil { + var changeDesc []string + if len(newFields) > 0 { + var names []string + for _, f := range newFields { + names = append(names, f.Name) + } + changeDesc = append(changeDesc, fmt.Sprintf("new columns %v", names)) + } + if len(updatedRecords) > 0 { + var names []string + for _, f := range updatedRecords { + names = append(names, f.Name) + } + changeDesc = append(changeDesc, fmt.Sprintf("updated RECORD fields %v", names)) + } + logger.Printf("Auto-upgraded table schema: %s", strings.Join(changeDesc, ", ")) + } + } else { + // No schema upgrade needed, just update label if missing or outdated + if storedVersion != schemaVersion { + update := bq.TableMetadataToUpdate{} + update.SetLabel(schemaVersionLabelKey, schemaVersion) + + _, err := existingTable.Update(ctx, update, meta.ETag) + if err != nil { + return fmt.Errorf("failed to update table labels: %w", err) + } + } + } + return nil +} + +func createAnalyticsViews(ctx context.Context, client *bq.Client, config Config) error { + for eventType, extraCols := range eventViewDefs { + viewName := fmt.Sprintf("%s_%s", config.ViewPrefix, strings.ToLower(eventType)) + + allCols := make([]string, 0, len(viewCommonColumns)+len(extraCols)) + allCols = append(allCols, viewCommonColumns...) + allCols = append(allCols, extraCols...) + columnsStr := strings.Join(allCols, ",\n ") + + sql := fmt.Sprintf( + viewSQLTemplate, + config.ProjectID, config.DatasetID, viewName, + columnsStr, + config.ProjectID, config.DatasetID, config.TableName, + eventType, + ) + + q := client.Query(sql) + job, err := q.Run(ctx) + if err != nil { + return fmt.Errorf("failed to run view creation query for %s: %w", viewName, err) + } + status, err := job.Wait(ctx) + if err != nil { + return fmt.Errorf("view creation query failed for %s: %w", viewName, err) + } + if err := status.Err(); err != nil { + return fmt.Errorf("view creation query execution failed for %s: %w", viewName, err) + } + } + return nil +} diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go index 2a59eefcf..4a8a79a0a 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" @@ -205,6 +206,8 @@ func setupTestPlugin(t *testing.T) (*baseplugin.Plugin, chan *storagepb.AppendRo config.ProjectID = "test-project" config.DatasetID = "test-dataset" config.TableName = "test-table" + config.CreateViews = false + config.AutoSchemaUpgrade = false mockTransport := &mockTransport{ roundTrip: func(r *http.Request) (*http.Response, error) { @@ -342,3 +345,405 @@ 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" + config.CreateViews = false + config.AutoSchemaUpgrade = false + + 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) + } +} + +func TestSchemaFieldsMatch(t *testing.T) { + existing := bq.Schema{ + {Name: "col1", Type: bq.StringFieldType}, + { + Name: "col_record", + Type: bq.RecordFieldType, + Schema: bq.Schema{ + {Name: "sub1", Type: bq.IntegerFieldType}, + }, + }, + } + + desired := bq.Schema{ + {Name: "col1", Type: bq.StringFieldType}, + {Name: "col2", Type: bq.BooleanFieldType}, // New top-level col + { + Name: "col_record", + Type: bq.RecordFieldType, + Schema: bq.Schema{ + {Name: "sub1", Type: bq.IntegerFieldType}, + {Name: "sub2", Type: bq.StringFieldType}, // New sub-field + }, + }, + } + + newFields, updatedRecords := schemaFieldsMatch(existing, desired) + + if len(newFields) != 1 || newFields[0].Name != "col2" { + t.Errorf("Expected 1 new top-level field 'col2', got %v", newFields) + } + + if len(updatedRecords) != 1 || updatedRecords[0].Name != "col_record" { + t.Fatalf("Expected 1 updated RECORD field 'col_record', got %v", updatedRecords) + } + + updatedRecord := updatedRecords[0] + if len(updatedRecord.Schema) != 2 { + t.Errorf("Expected updated RECORD to have 2 sub-fields, got %d", len(updatedRecord.Schema)) + } + if updatedRecord.Schema[1].Name != "sub2" { + t.Errorf("Expected second sub-field of RECORD to be 'sub2', got %s", updatedRecord.Schema[1].Name) + } +} + +func TestNewBigQueryAgentAnalyticsPlugin_TableExists(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + config.AutoSchemaUpgrade = false + config.CreateViews = false + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"tableReference":{"projectId":"test-project","datasetId":"test-dataset","tableId":"test-table"}}`)), + }, 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) + } + + p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err != nil { + t.Fatalf("Expected no error when table exists, got: %v", err) + } + if p == nil { + t.Fatal("Expected plugin to be non-nil") + } +} + +func TestNewBigQueryAgentAnalyticsPlugin_TableNotFound_CreateSucceeds(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + config.AutoSchemaUpgrade = false + config.CreateViews = false + + createCalled := false + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + 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 + } + if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") { + createCalled = true + 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) + } + + p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err != nil { + t.Fatalf("Expected no error when table is created successfully, got: %v", err) + } + if p == nil { + t.Fatal("Expected plugin to be non-nil") + } + if !createCalled { + t.Error("Expected table creation to be attempted, but it was not called") + } +} + +func TestNewBigQueryAgentAnalyticsPlugin_TableNotFound_CreateFails(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + config.AutoSchemaUpgrade = false + config.CreateViews = false + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + 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 + } + if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") { + return &http.Response{ + StatusCode: http.StatusForbidden, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":403,"message":"Permission denied"}}`)), + }, 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) + } + + p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err == nil { + t.Fatal("Expected error when table creation fails, got nil") + } + if p != nil { + t.Error("Expected plugin to be nil when table creation fails") + } + if !strings.Contains(err.Error(), "failed to create BigQuery table") { + t.Errorf("Expected error message to contain 'failed to create BigQuery table', got: %v", err) + } +} + +func TestNewBigQueryAgentAnalyticsPlugin_MetadataError_Non404(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + config.AutoSchemaUpgrade = false + config.CreateViews = false + + createAttempted := false + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") { + return &http.Response{ + StatusCode: http.StatusForbidden, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":403,"message":"Permission denied"}}`)), + }, nil + } + if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") { + createAttempted = true + 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) + } + + p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err == nil { + t.Fatal("Expected error when metadata retrieval fails with non-404 error, got nil") + } + if p != nil { + t.Error("Expected plugin to be nil when metadata retrieval fails") + } + if createAttempted { + t.Error("Expected table creation NOT to be attempted, but it was called") + } + if !strings.Contains(err.Error(), "failed to retrieve BigQuery table metadata") { + t.Errorf("Expected error message to contain 'failed to retrieve BigQuery table metadata', got: %v", err) + } +} diff --git a/plugin/agentanalytics/config.go b/plugin/agentanalytics/config.go index 40e3b5595..2ae58df05 100644 --- a/plugin/agentanalytics/config.go +++ b/plugin/agentanalytics/config.go @@ -63,6 +63,15 @@ type Config struct { // Retry configuration for appending rows. RetryConfig RetryConfig + + // Automatically add new columns to existing tables when schema evolves. + AutoSchemaUpgrade bool + + // Automatically create per-event-type BigQuery views. + CreateViews bool + + // Prefix for auto-created view names. + ViewPrefix string } // DefaultConfig returns the default configuration for the agent analytics plugin. @@ -86,5 +95,8 @@ func DefaultConfig() Config { MaxDelay: 10 * time.Second, Multiplier: 2.0, }, + AutoSchemaUpgrade: true, + CreateViews: true, + ViewPrefix: "v", } }