diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/README.md b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/README.md index ee8691df..2c1c5208 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/README.md +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/README.md @@ -31,7 +31,6 @@ see the [quick start](#quick-start). - [Unit Tests](#unit-tests) - [Compliance Tests](#compliance-tests) - [Fuzzing Tests](#fuzzing-tests) -- [Generating Mock Files Using Mockery](./mocks/README.md) - [Connection Methods](#connection-methods) - [Plain TCP Connection](#plain-tcp-connection) - [Secured TCP Connection (TLS)](#secured-tcp-connection-tls) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable.go index 85476a6e..f7a4259b 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable.go @@ -22,6 +22,7 @@ import ( "encoding/hex" "fmt" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "go.opentelemetry.io/otel/trace" "strings" "time" @@ -31,7 +32,6 @@ import ( btpb "cloud.google.com/go/bigtable/apiv2/bigtablepb" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata" - otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" "github.com/datastax/go-cassandra-native-protocol/message" "go.uber.org/zap" ) @@ -56,17 +56,17 @@ func NewBigtableClient(clients *types.BigtableClientManager, logger *zap.Logger, func (btc *BigtableAdapter) Execute(ctx context.Context, query types.IExecutableQuery) (message.Message, error) { switch q := query.(type) { case *types.BoundDeleteQuery: - return btc.DeleteRow(ctx, q) + return btc.deleteRow(ctx, q) case *types.BigtableWriteMutation: return btc.mutateRow(ctx, q) case *types.ExecutableSelectQuery: - return btc.ExecutePreparedStatement(ctx, q) + return btc.executePreparedStatement(ctx, q) case *types.CreateTableStatementMap: return btc.schemaManager.CreateTable(ctx, q) case *types.AlterTableStatementMap: return btc.schemaManager.AlterTable(ctx, q) case *types.TruncateTableStatementMap: - err := btc.DropAllRows(ctx, q) + err := btc.dropAllRows(ctx, q) return emptyRowsResult(), err case *types.DropTableQuery: return btc.schemaManager.DropTable(ctx, q) @@ -88,11 +88,9 @@ func (btc *BigtableAdapter) Execute(ctx context.Context, query types.IExecutable // Returns: // - error: Error if the mutation fails. func (btc *BigtableAdapter) mutateRow(ctx context.Context, input *types.BigtableWriteMutation) (message.Message, error) { - otelgo.AddAnnotation(ctx, applyingBigtableMutation) + span := trace.SpanFromContext(ctx) mut := bigtable.NewMutation() - btc.Logger.Info("mutating row", zap.String("key", hex.EncodeToString([]byte(input.RowKey())))) - client, err := btc.clients.GetClient(input.Keyspace()) if err != nil { return nil, err @@ -119,7 +117,7 @@ func (btc *BigtableAdapter) mutateRow(ctx context.Context, input *types.Bigtable } err := tbl.Apply(ctx, string(input.RowKey()), conditionalMutation, bigtable.GetCondMutationResult(&matched)) - otelgo.AddAnnotation(ctx, bigtableMutationApplied) + span.AddEvent(bigtableMutationApplied) if err != nil { return nil, err } @@ -129,7 +127,7 @@ func (btc *BigtableAdapter) mutateRow(ctx context.Context, input *types.Bigtable // If no conditions, apply the mutation directly err = tbl.Apply(ctx, string(input.RowKey()), mut) - otelgo.AddAnnotation(ctx, bigtableMutationApplied) + span.AddEvent(bigtableMutationApplied) if err != nil { return nil, err } @@ -185,13 +183,13 @@ func (btc *BigtableAdapter) buildMutation(ctx context.Context, table *bigtable.T return nil } -func (btc *BigtableAdapter) DropAllRows(ctx context.Context, data *types.TruncateTableStatementMap) error { +func (btc *BigtableAdapter) dropAllRows(ctx context.Context, data *types.TruncateTableStatementMap) error { _, err := btc.schemaManager.Schemas().GetTableSchema(data.Keyspace(), data.Table()) if err != nil { return err } - // performance optimization because DropAllRows can be slow + // performance optimization because dropAllRows can be slow hasRows, err := btc.hasAnyRows(ctx, data.Keyspace(), data.Table()) if err != nil { return err @@ -263,7 +261,7 @@ func (btc *BigtableAdapter) InsertRow(ctx context.Context, input *types.Bigtable return btc.mutateRow(ctx, input) } -// UpdateRow - Updates a row in the specified bigtable table. +// updateRow - Updates a row in the specified bigtable table. // // Parameters: // - ctx: Context for the operation, used for cancellation and deadlines. @@ -271,12 +269,13 @@ func (btc *BigtableAdapter) InsertRow(ctx context.Context, input *types.Bigtable // // Returns: // - error: Error if the update fails. -func (btc *BigtableAdapter) UpdateRow(ctx context.Context, input *types.BigtableWriteMutation) (message.Message, error) { +func (btc *BigtableAdapter) updateRow(ctx context.Context, input *types.BigtableWriteMutation) (message.Message, error) { return btc.mutateRow(ctx, input) } -func (btc *BigtableAdapter) DeleteRow(ctx context.Context, deleteQueryData *types.BoundDeleteQuery) (message.Message, error) { - otelgo.AddAnnotation(ctx, applyingDeleteMutation) +func (btc *BigtableAdapter) deleteRow(ctx context.Context, deleteQueryData *types.BoundDeleteQuery) (message.Message, error) { + span := trace.SpanFromContext(ctx) + span.AddEvent(applyingDeleteMutation) client, err := btc.clients.GetClient(deleteQueryData.Keyspace()) if err != nil { return nil, err @@ -308,7 +307,7 @@ func (btc *BigtableAdapter) DeleteRow(ctx context.Context, deleteQueryData *type return nil, err } } - otelgo.AddAnnotation(ctx, deleteMutationApplied) + span.AddEvent(deleteMutationApplied) return &message.VoidResult{}, nil } @@ -345,6 +344,7 @@ func (btc *BigtableAdapter) buildDeleteMutation(ctx context.Context, table *bigt // - BulkOperationResponse: Response indicating the result of the bulk operation. // - error: Error if the bulk mutation fails. func (btc *BigtableAdapter) ApplyBulkMutation(ctx context.Context, keyspace types.Keyspace, tableName types.TableName, mutationData []types.IBigtableMutation) (BulkOperationResponse, error) { + span := trace.SpanFromContext(ctx) client, err := btc.clients.GetClient(keyspace) if err != nil { return BulkOperationResponse{ @@ -385,9 +385,10 @@ func (btc *BigtableAdapter) ApplyBulkMutation(ctx context.Context, keyspace type }, err } default: + err := fmt.Errorf("unhandled bulk mutation type %T", md) return BulkOperationResponse{ FailedRows: fmt.Sprintf("All Rows are failed because: unsupported bulk operation: %T", v), - }, fmt.Errorf("unhandled bulk mutation type %T", md) + }, err } } // create mutations from mutation data @@ -398,7 +399,7 @@ func (btc *BigtableAdapter) ApplyBulkMutation(ctx context.Context, keyspace type mutations = append(mutations, mutation) rowKeys = append(rowKeys, string(key)) } - otelgo.AddAnnotation(ctx, applyingBulkMutation) + span.AddEvent(applyingBulkMutation) errs, err := table.ApplyBulk(ctx, rowKeys, mutations) if err != nil { @@ -423,7 +424,7 @@ func (btc *BigtableAdapter) ApplyBulkMutation(ctx context.Context, keyspace type FailedRows: "", } } - otelgo.AddAnnotation(ctx, bulkMutationApplied) + span.AddEvent(bulkMutationApplied) return res, nil } @@ -510,8 +511,8 @@ func (btc *BigtableAdapter) PrepareStatement(ctx context.Context, query types.IP return nil, nil } - selectQuery, ok := query.(*types.PreparedSelectQuery) - if !ok { + selectQuery, isType := query.(*types.PreparedSelectQuery) + if !isType { // only select queries can be prepared in Bigtable at this time return nil, nil } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_select.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_select.go index e08a7a9d..1fe4e194 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_select.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_select.go @@ -22,6 +22,7 @@ import ( "fmt" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/constants" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/utilities" @@ -32,7 +33,7 @@ import ( "time" ) -// ExecutePreparedStatement - Executes a prepared statement on Bigtable and returns the result. +// executePreparedStatement - Executes a prepared statement on Bigtable and returns the result. // Parameters: // - ctx: Context for the operation, used for cancellation and deadlines. // - query: rh.QueryMetadata containing the query and parameters. @@ -42,7 +43,7 @@ import ( // - *message.RowsResult: The result of the select statement. // - time.Duration: The total elapsed time for the operation. // - error: Error if the statement preparation or execution fails. -func (btc *BigtableAdapter) ExecutePreparedStatement(ctx context.Context, query *types.ExecutableSelectQuery) (*message.RowsResult, error) { +func (btc *BigtableAdapter) executePreparedStatement(ctx context.Context, query *types.ExecutableSelectQuery) (*message.RowsResult, error) { if query.CachedBTPrepare == nil { return nil, fmt.Errorf("cannot execute select query because prepared bigtable query is nil") } @@ -59,14 +60,19 @@ func (btc *BigtableAdapter) ExecutePreparedStatement(ctx context.Context, query return nil, fmt.Errorf("failed to bind parameters: %w", err) } + table, err := btc.schemaManager.Schemas().GetTableSchema(query.Keyspace(), query.Table()) + if err != nil { + return nil, err + } + var processingErr error var rows []types.GoRow executeErr := boundStmt.Execute(ctx, func(resultRow bigtable.ResultRow) bool { - r, convertErr := btc.convertResultRow(resultRow, query) // Call the implemented helper + r, convertErr := btc.convertResultRow(resultRow, query, table) if convertErr != nil { btc.Logger.Error("Failed to convert result row", zap.Error(convertErr), zap.String("btql", query.TranslatedQuery)) - processingErr = convertErr // Capture the error - return false // Stop execution + processingErr = convertErr + return false // Stop execution } rows = append(rows, r) return true // Continue processing @@ -82,12 +88,7 @@ func (btc *BigtableAdapter) ExecutePreparedStatement(ctx context.Context, query return responsehandler.BuildRowsResultResponse(query, rows, query.ProtocolVersion) } -func (btc *BigtableAdapter) convertResultRow(resultRow bigtable.ResultRow, query *types.ExecutableSelectQuery) (types.GoRow, error) { - table, err := btc.schemaManager.Schemas().GetTableSchema(query.Keyspace(), query.Table()) - if err != nil { - return nil, err - } - +func (btc *BigtableAdapter) convertResultRow(resultRow bigtable.ResultRow, query *types.ExecutableSelectQuery, table *metadata.TableSchema) (types.GoRow, error) { result := make(types.GoRow) for i, colMeta := range resultRow.Metadata.Columns { var val any @@ -140,10 +141,6 @@ func (btc *BigtableAdapter) convertResultRow(resultRow bigtable.ResultRow, query return nil, fmt.Errorf("result already set for column `%s`", key) } - if key == "list_text" { - btc.Logger.Log(zap.InfoLevel, "list_text", zap.Any("value", val)) - } - goValue, err := rowValueToGoValue(val, expectedType) if err != nil { return nil, fmt.Errorf("failed to convert result for '%s': %w", key, err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_test.go index 12b9e812..66dd8e46 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable/bigtable_test.go @@ -166,7 +166,7 @@ func TestDeleteRow(t *testing.T) { require.NoError(t, err) deleteQuery := types.NewBoundDeleteQuery(keyspace, tableName, "", rowKey, false, nil) - _, err = btc.DeleteRow(t.Context(), deleteQuery) + _, err = btc.deleteRow(t.Context(), deleteQuery) require.NoError(t, err) // Verify deletion @@ -236,7 +236,7 @@ func TestMutateRowDeleteColumnFamily(t *testing.T) { // Delete cf2 updateData := types.NewBigtableWriteMutation(keyspace, tableName, "", types.IfSpec{}, types.QueryTypeUpdate, key) updateData.AddMutations(types.NewDeleteCellsOp("tags")) - _, err = btc.UpdateRow(t.Context(), updateData) + _, err = btc.updateRow(t.Context(), updateData) require.NoError(t, err) // Verify deletion by reading the row @@ -259,7 +259,7 @@ func TestMutateRowDeleteQualifiers(t *testing.T) { // Delete col1 updateData := types.NewBigtableWriteMutation(keyspace, tableName, "", types.IfSpec{}, types.QueryTypeUpdate, key) updateData.AddMutations(types.NewDeleteColumnOp(types.BigtableColumn{Family: "cf1", Column: "col1"})) - _, err = btc.UpdateRow(t.Context(), updateData) + _, err = btc.updateRow(t.Context(), updateData) require.NoError(t, err) // Verify deletion by reading the row @@ -289,7 +289,7 @@ func TestMutateRowIfExists(t *testing.T) { // Update the row when it exists updateData := types.NewBigtableWriteMutation(keyspace, tableName, "", types.IfSpec{IfExists: true}, types.QueryTypeUpdate, key1) updateData.AddMutations(types.NewWriteCellOp("cf1", "col1", []byte("v2"))) - res, err := btc.UpdateRow(t.Context(), updateData) + res, err := btc.updateRow(t.Context(), updateData) require.NoError(t, err) assert.True(t, wasApplied(res)) @@ -305,7 +305,7 @@ func TestMutateRowIfExists(t *testing.T) { // Attempt to update a non-existent row updateDataNonExistent := types.NewBigtableWriteMutation(keyspace, tableName, "", types.IfSpec{IfExists: true}, types.QueryTypeUpdate, key2) updateDataNonExistent.AddMutations(types.NewWriteCellOp("cf1", "col1", []byte("v2"))) - res, err = btc.UpdateRow(t.Context(), updateDataNonExistent) + res, err = btc.updateRow(t.Context(), updateDataNonExistent) require.NoError(t, err) assert.False(t, wasApplied(res)) @@ -350,7 +350,7 @@ func TestMutateRowInvalidKeyspace(t *testing.T) { updateData := types.NewBigtableWriteMutation("invalid-keyspace", "any-table", "", types.IfSpec{}, types.QueryTypeUpdate, "row1") updateData.AddMutations(types.NewWriteCellOp("cf1", "col1", []byte("value"))) - _, err := localBtc.UpdateRow(t.Context(), updateData) + _, err := localBtc.updateRow(t.Context(), updateData) require.Error(t, err) assert.Contains(t, err.Error(), "bigtable client not found for keyspace 'invalid-keyspace'") } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/config.yaml b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/config.yaml index c61d65f8..68c221f4 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/config.yaml +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/config.yaml @@ -92,6 +92,9 @@ otel: metrics: # Collector service endpoint endpoint: METRICS_SERVICE_ENDPOINT + # Set to true to export metrics directly to Google Cloud Monitoring. + # If true, 'endpoint' will be ignored, and metrics will be sent to the project specified in traces.projectId. + gcpMetricsEnabled: False traces: # Collector service endpoint diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/example_config.yaml b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/example_config.yaml index 6dcb58ea..8c01b298 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/example_config.yaml +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/example_config.yaml @@ -33,4 +33,27 @@ listeners: grpcChannels: 4 otel: + # Set enabled to true or false for OTEL metrics/traces/logs. enabled: False + # Name of the collector service to be setup as a sidecar + serviceName: YOUR_OTEL_COLLECTOR_SERVICE_NAME + + healthcheck: + # Enable/Disable Health Check for OTEL, Default 'False'. + enabled: False + # Health check endpoint for the OTEL collector service + endpoint: localhost:13133 + metrics: + # Collector service endpoint + endpoint: localhost:4317 + # Set to true to export metrics directly to Google Cloud Monitoring. + # If true, 'endpoint' will be ignored, and metrics will be sent to the project specified in traces.projectId. + gcpMetricsEnabled: False + + traces: + # Collector service endpoint + endpoint: localhost:4317 + # Project ID to use for exporting traces to Google Cloud Trace. + projectId: YOUR_GCP_PROJECT_ID + #Sampling ratio should be between 0 and 1. Here 0.05 means 5/100 Sampling ratio. + samplingRatio: 1 diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors/executor_manager.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors/executor_manager.go index 160d9dbd..28ce6983 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors/executor_manager.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors/executor_manager.go @@ -7,7 +7,11 @@ import ( "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/mem_table" schemaMapping "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata" + otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" "github.com/datastax/go-cassandra-native-protocol/message" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" "go.uber.org/zap" "strings" ) @@ -20,9 +24,10 @@ type IQueryExecutor interface { type QueryExecutorManager struct { logger *zap.Logger executors []IQueryExecutor + trace trace.Tracer } -func NewQueryExecutorManager(logger *zap.Logger, s *schemaMapping.SchemaMetadata, bt *bigtableModule.BigtableAdapter, systemTables *mem_table.InMemEngine) *QueryExecutorManager { +func NewQueryExecutorManager(logger *zap.Logger, s *schemaMapping.SchemaMetadata, bt *bigtableModule.BigtableAdapter, systemTables *mem_table.InMemEngine, otelInst *otelgo.OpenTelemetry) *QueryExecutorManager { return &QueryExecutorManager{ logger: logger, executors: []IQueryExecutor{ @@ -31,15 +36,37 @@ func NewQueryExecutorManager(logger *zap.Logger, s *schemaMapping.SchemaMetadata newSelectSystemTableExecutor(s, systemTables), newBigtableExecutor(bt), }, + trace: otel.GetTracerProvider().Tracer("executor"), } } -func (m *QueryExecutorManager) Execute(ctx context.Context, client types.ICassandraClient, q types.IExecutableQuery) (message.Message, error) { +func (m *QueryExecutorManager) getExecutor(q types.IExecutableQuery) (IQueryExecutor, error) { for _, e := range m.executors { if e.CanRun(q) { - m.logger.Debug("executing query", zap.String("cql", q.CqlQuery()), zap.String("btql", q.BigtableQuery())) - return e.Execute(ctx, client, q) + return e, nil } } return nil, fmt.Errorf("no executor found for query %s on keyspace %s", strings.ToUpper(q.QueryType().String()), q.Keyspace()) } + +func (m *QueryExecutorManager) Execute(ctx context.Context, client types.ICassandraClient, q types.IExecutableQuery) (message.Message, error) { + otelCtx, span := m.trace.Start(ctx, "execute") + defer span.End() + + otelgo.AddQueryAnnotations(span, q) + + executor, err := m.getExecutor(q) + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return nil, err + } + + msg, err := executor.Execute(otelCtx, client, q) + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return nil, err + } + return msg, nil +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/attributes.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/attributes.go new file mode 100644 index 00000000..3b2cb70e --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/attributes.go @@ -0,0 +1,16 @@ +package types + +type Attributes struct { + Method string + Err error + QueryType QueryType + Keyspace Keyspace + Table TableName + Status string +} + +func (a *Attributes) SetQueryInfo(q IQuery) { + a.Keyspace = q.Keyspace() + a.Table = q.Table() + a.QueryType = q.QueryType() +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/bigtable_client_manager.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/bigtable_client_manager.go index bd862554..d4adf8e0 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/bigtable_client_manager.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/bigtable_client_manager.go @@ -5,7 +5,8 @@ import ( "context" "fmt" "google.golang.org/api/option" - "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) type BigtableClientSet struct { @@ -42,12 +43,27 @@ func createBigtableClientSet(ctx context.Context, config *ProxyInstanceConfig, i if err != nil { return nil, fmt.Errorf("failed to create admin client for keyspace `%s`: %v", instanceMapping.Keyspace, err) } - // Create a gRPC connection pool with the specified number of channels - pool := grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(1024 * 1024 * 10)) // 10 MB max message size + + err = adminClient.CreateTable(ctx, string(config.BigtableConfig.SchemaMappingTable)) + if status.Code(err) == codes.AlreadyExists { + // continue - maybe another Proxy instance raced, and created it instead + err = nil + } + if err != nil { + return nil, err + } + + // todo clean up + err = adminClient.CreateColumnFamily(ctx, string(config.BigtableConfig.SchemaMappingTable), "cf") + if status.Code(err) == codes.AlreadyExists { + err = nil + } + if err != nil { + return nil, err + } // Specify gRPC connection options, including your custom number of channels opts := []option.ClientOption{ - option.WithGRPCDialOption(pool), option.WithGRPCConnectionPool(config.BigtableConfig.Session.GrpcChannels), option.WithUserAgent(config.Options.UserAgent), } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/config_types.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/config_types.go index 78cc68cf..531f713b 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/config_types.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/config_types.go @@ -36,16 +36,19 @@ type CliArgs struct { } type OtelConfig struct { - Enabled bool - ServiceName string - HealthCheck struct { + Enabled bool + ServiceName string + ServiceVersion string + HealthCheck struct { Enabled bool Endpoint string } Metrics struct { - Endpoint string + Endpoint string + GcpMetricsEnabled bool } Traces struct { + ProjectId string Endpoint string SamplingRatio float64 } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/types.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/types.go index 333a4320..e2059d98 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/types.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types/types.go @@ -156,7 +156,6 @@ const ( QueryTypeInsert QueryTypeUpdate QueryTypeDelete - // ddl QueryTypeCreate QueryTypeAlter QueryTypeDrop @@ -191,20 +190,28 @@ func (q QueryType) String() string { } func (q QueryType) IsDDLType() bool { switch q { - case QueryTypeTruncate, QueryTypeDrop, QueryTypeAlter, QueryTypeCreate: + case QueryTypeDrop, QueryTypeAlter, QueryTypeCreate: return true default: return false } } +type IQuery interface { + Keyspace() Keyspace + Table() TableName + QueryType() QueryType +} + type IExecutableQuery interface { Keyspace() Keyspace + Table() TableName // empty string if no table involved e.g. "USE keyspace;" QueryType() QueryType AsBulkMutation() (IBigtableMutation, bool) CqlQuery() string BigtableQuery() string } + type IPreparedQuery interface { Keyspace() Keyspace Table() TableName // empty string if no table involved e.g. "USE keyspace;" @@ -229,8 +236,9 @@ type RawQuery struct { cql string qt QueryType sessionKeyspace Keyspace - parser *parser.ProxyCqlParser - startTime time.Time + // warning: parsers are pooled for performance reasons. this will be released back into the pool and set to nil after translation + parser *parser.ProxyCqlParser + startTime time.Time } func NewRawQuery(header *frame.Header, sessionKeyspace Keyspace, cql string, parser *parser.ProxyCqlParser, qt QueryType) *RawQuery { @@ -248,6 +256,11 @@ func NewRawQueryWithTime(header *frame.Header, sessionKeyspace Keyspace, cql str } } +func (r *RawQuery) Release() { + r.parser.Release() + r.parser = nil +} + func (r *RawQuery) QueryType() QueryType { return r.qt } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.mod b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.mod index 85346c4f..929ef481 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.mod +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.mod @@ -1,88 +1,90 @@ module github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy -go 1.24.0 +go 1.25.8 require ( + github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.57.0 + github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.33.0 github.com/alecthomas/kong v0.2.17 github.com/antlr4-go/antlr/v4 v4.13.0 github.com/datastax/go-cassandra-native-protocol v0.0.0-20211124104234-f6aea54fa801 github.com/gocql/gocql v1.7.0 - github.com/google/go-cmp v0.7.0 github.com/google/uuid v1.6.0 github.com/hashicorp/golang-lru v0.5.4 - github.com/stretchr/testify v1.10.0 - github.com/tj/assert v0.0.3 - go.opentelemetry.io/contrib/detectors/gcp v1.34.0 - go.opentelemetry.io/otel v1.35.0 + github.com/stretchr/testify v1.11.1 + go.opentelemetry.io/contrib/detectors/gcp v1.42.0 + go.opentelemetry.io/otel v1.44.0 go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.24.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0 - go.opentelemetry.io/otel/metric v1.35.0 - go.opentelemetry.io/otel/sdk v1.35.0 - go.opentelemetry.io/otel/sdk/metric v1.35.0 - go.opentelemetry.io/otel/trace v1.35.0 - go.uber.org/atomic v1.11.0 + go.opentelemetry.io/otel/metric v1.44.0 + go.opentelemetry.io/otel/sdk v1.44.0 + go.opentelemetry.io/otel/sdk/metric v1.44.0 + go.opentelemetry.io/otel/trace v1.44.0 go.uber.org/zap v1.26.0 - google.golang.org/api v0.228.0 - google.golang.org/grpc v1.71.1 + google.golang.org/api v0.283.0 + google.golang.org/grpc v1.81.1 gopkg.in/yaml.v2 v2.4.0 ) require ( - cel.dev/expr v0.23.0 // indirect - cloud.google.com/go v0.120.0 // indirect - cloud.google.com/go/auth v0.15.0 // indirect + cel.dev/expr v0.25.1 // indirect + cloud.google.com/go v0.123.0 // indirect + cloud.google.com/go/auth v0.20.0 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect - cloud.google.com/go/iam v1.5.0 // indirect - cloud.google.com/go/longrunning v0.6.6 // indirect - cloud.google.com/go/monitoring v1.24.1 // indirect + cloud.google.com/go/iam v1.11.0 // indirect + cloud.google.com/go/longrunning v1.0.0 // indirect + cloud.google.com/go/monitoring v1.29.0 // indirect + cloud.google.com/go/trace v1.16.0 // indirect github.com/BurntSushi/toml v1.5.0 // indirect - github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.25.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.31.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.57.0 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect - github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect + github.com/envoyproxy/go-control-plane/envoy v1.37.0 // indirect + github.com/go-jose/go-jose/v4 v4.1.4 // indirect github.com/golang/snappy v0.0.4 // indirect github.com/google/btree v1.1.3 // indirect - github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 // indirect github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed // indirect github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect - github.com/stretchr/objx v0.5.2 // indirect - go.opentelemetry.io/auto/sdk v1.1.0 // indirect + github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.31.0 // indirect - go.opentelemetry.io/proto/otlp v1.3.1 // indirect + go.opentelemetry.io/proto/otlp v1.9.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect rsc.io/binaryregexp v0.2.0 // indirect ) require ( - cloud.google.com/go/bigtable v1.36.0 - cloud.google.com/go/compute/metadata v0.6.0 // indirect + cloud.google.com/go/bigtable v1.49.0 + cloud.google.com/go/compute/metadata v0.9.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f // indirect - github.com/davecgh/go-spew v1.1.1 // indirect - github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect + github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/envoyproxy/protoc-gen-validate v1.3.3 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect - github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/google/s2a-go v0.1.9 // indirect - github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect - github.com/googleapis/gax-go/v2 v2.14.1 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.16 // indirect + github.com/googleapis/gax-go/v2 v2.22.0 // indirect github.com/natefinch/lumberjack v2.0.0+incompatible github.com/pkg/errors v0.9.1 // indirect - github.com/pmezard/go-difflib v1.0.0 // indirect - go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 // indirect - go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.68.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0 // indirect go.uber.org/multierr v1.11.0 // indirect - golang.org/x/crypto v0.36.0 // indirect + golang.org/x/crypto v0.51.0 // indirect golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63 - golang.org/x/net v0.38.0 // indirect - golang.org/x/oauth2 v0.28.0 // indirect - golang.org/x/sync v0.12.0 // indirect - golang.org/x/sys v0.31.0 - golang.org/x/text v0.23.0 // indirect - golang.org/x/time v0.11.0 // indirect - google.golang.org/genproto v0.0.0-20250404141209-ee84b53bf3d0 // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20250404141209-ee84b53bf3d0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20250404141209-ee84b53bf3d0 // indirect - google.golang.org/protobuf v1.36.6 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/oauth2 v0.36.0 // indirect + golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.45.0 + golang.org/x/text v0.37.0 // indirect + golang.org/x/time v0.15.0 // indirect + google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect + google.golang.org/protobuf v1.36.11 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.sum b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.sum index f57b117d..eda1deb2 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.sum +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/go.sum @@ -1,25 +1,37 @@ -cel.dev/expr v0.23.0 h1:wUb94w6OYQS4uXraxo9U+wUAs9jT47Xvl4iPgAwM2ss= -cel.dev/expr v0.23.0/go.mod h1:hLPLo1W4QUmuYdA72RBX06QTs6MXw941piREPl3Yfiw= -cloud.google.com/go v0.120.0 h1:wc6bgG9DHyKqF5/vQvX1CiZrtHnxJjBlKUyF9nP6meA= -cloud.google.com/go v0.120.0/go.mod h1:/beW32s8/pGRuj4IILWQNd4uuebeT4dkOhKmkfit64Q= -cloud.google.com/go/auth v0.15.0 h1:Ly0u4aA5vG/fsSsxu98qCQBemXtAtJf+95z9HK+cxps= -cloud.google.com/go/auth v0.15.0/go.mod h1:WJDGqZ1o9E9wKIL+IwStfyn/+s59zl4Bi+1KQNVXLZ8= +cel.dev/expr v0.25.1 h1:1KrZg61W6TWSxuNZ37Xy49ps13NUovb66QLprthtwi4= +cel.dev/expr v0.25.1/go.mod h1:hrXvqGP6G6gyx8UAHSHJ5RGk//1Oj5nXQ2NI02Nrsg4= +cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= +cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= +cloud.google.com/go/auth v0.20.0 h1:kXTssoVb4azsVDoUiF8KvxAqrsQcQtB53DcSgta74CA= +cloud.google.com/go/auth v0.20.0/go.mod h1:942/yi/itH1SsmpyrbnTMDgGfdy2BUqIKyd0cyYLc5Q= cloud.google.com/go/auth/oauth2adapt v0.2.8 h1:keo8NaayQZ6wimpNSmW5OPc283g65QNIiLpZnkHRbnc= cloud.google.com/go/auth/oauth2adapt v0.2.8/go.mod h1:XQ9y31RkqZCcwJWNSx2Xvric3RrU88hAYYbjDWYDL+c= -cloud.google.com/go/bigtable v1.36.0 h1:GU4XWYb7H9XYHvksDA/jixUZUv2ZESNS48QEc/qduV4= -cloud.google.com/go/bigtable v1.36.0/go.mod h1:u98oqNAXiAufepkRGAd95lq2ap4kHGr3wLeFojvJwew= -cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I= -cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg= -cloud.google.com/go/iam v1.5.0 h1:QlLcVMhbLGOjRcGe6VTGGTyQib8dRLK2B/kYNV0+2xs= -cloud.google.com/go/iam v1.5.0/go.mod h1:U+DOtKQltF/LxPEtcDLoobcsZMilSRwR7mgNL7knOpo= -cloud.google.com/go/longrunning v0.6.6 h1:XJNDo5MUfMM05xK3ewpbSdmt7R2Zw+aQEMbdQR65Rbw= -cloud.google.com/go/longrunning v0.6.6/go.mod h1:hyeGJUrPHcx0u2Uu1UFSoYZLn4lkMrccJig0t4FI7yw= -cloud.google.com/go/monitoring v1.24.1 h1:vKiypZVFD/5a3BbQMvI4gZdl8445ITzXFh257XBgrS0= -cloud.google.com/go/monitoring v1.24.1/go.mod h1:Z05d1/vn9NaujqY2voG6pVQXoJGbp+r3laV+LySt9K0= +cloud.google.com/go/bigtable v1.49.0 h1:kdA1jbO8EJIMI752zD50o5z6Wu82cvCwIibhrpJK0tI= +cloud.google.com/go/bigtable v1.49.0/go.mod h1:RTannV5mvoJM8KscLTfRYMPo84u9/j+C3PSyYJGf5Ic= +cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= +cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= +cloud.google.com/go/iam v1.11.0 h1:KieQ9Pb+LLPak1O3Rv3GgCxhnmkYf7Xyh0P5HfF1jFM= +cloud.google.com/go/iam v1.11.0/go.mod h1:KP+nKGugNJW4LcLx1uEZcq1ok5sQHFaQehQNl4QDgV4= +cloud.google.com/go/logging v1.18.0 h1:KhzZq+1cSkPH9YUaKLLhLtQxIHitVayBmk0sGfoM9+k= +cloud.google.com/go/logging v1.18.0/go.mod h1:ZGKnpBaURITh+g/uom2VhbiFoFWvejcrHPDhxFtU/gI= +cloud.google.com/go/longrunning v1.0.0 h1:lwzWEYD8+NkYV7dhexOz6kmlvajZA70+bW/xMhRVVdY= +cloud.google.com/go/longrunning v1.0.0/go.mod h1:8nqFBPOO1U/XkhWl0I19AMZEphrHi73VNABIpKYaTwM= +cloud.google.com/go/monitoring v1.29.0 h1:AHhDsFaSax1/4k+qlIDX/SDGe6hggnfXJ9dkgD9qBPY= +cloud.google.com/go/monitoring v1.29.0/go.mod h1:72NOVjJXHY/HBfoLT0+qlCZBT059+9VXLeAnL2PeeVM= +cloud.google.com/go/trace v1.16.0 h1:GmQovzFc5F0CNfl0VLgL64aoTtu7xsM0YajW2GlG9+E= +cloud.google.com/go/trace v1.16.0/go.mod h1:r+bdAn16dKLSV1G2D5v3e58IlQlizfxWrUfjx7kM7X0= github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg= github.com/BurntSushi/toml v1.5.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= -github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.25.0 h1:3c8yed4lgqTt+oTQ+JNMDo+F4xprBf+O/il4ZC0nRLw= -github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.25.0/go.mod h1:obipzmGjfSjam60XLwGfqUkJsfiheAl+TUjG+4yzyPM= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.31.0 h1:DHa2U07rk8syqvCge0QIGMCE1WxGj9njT44GH7zNJLQ= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.31.0/go.mod h1:P4WPRUkOhJC13W//jWpyfJNDAIpvRbAUIYLX/4jtlE0= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.57.0 h1:jLdiS1vO+XJFyDSWRHBx56r4s/NNtcl5J6KyCcWUX/w= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.57.0/go.mod h1:8lmpHY+1VRoteiOwyrQMDt1YGXOrFKCz+1wJW7n3ODY= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.33.0 h1:92kCbSANHUnbS7tzc9uhUx4H+MF1Lr/QXInoPkjz6JU= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.33.0/go.mod h1:n2lr2dT7IuNeGhW6iB/+nJB+Fo9vZxCGwEQOy2Hc15Y= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.57.0 h1:cSjUzZ7KU8hicTgzaSv9NmSyM9fTVK3y5lsBUl3wOis= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.57.0/go.mod h1:dzcEjy1WJ0Q4u9twNR3LcLhNoYMRCrMCMafpxa0TjPQ= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.57.0 h1:RoO5+d7uCmDqovLrHCr2/BuViUXvdcrNxyNM1pN9dDQ= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.57.0/go.mod h1:YqwkQPrWSC7+byyc1VlKbWLBF5JsW5IoL6xUkemYSXk= github.com/alecthomas/kong v0.2.17 h1:URDISCI96MIgcIlQyoCAlhOmrSw6pZScBNkctg8r0W0= github.com/alecthomas/kong v0.2.17/go.mod h1:ka3VZ8GZNPXv9Ov+j4YNLkI8mTuhXyr/0ktSlqIydQQ= github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI= @@ -32,27 +44,30 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f h1:C5bqEmzEPLsHm9Mv73lSE9e9bKV23aB1vxOsmZrkl3k= -github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 h1:aBangftG7EVZoUb69Os8IaYg++6uMOdKK83QtkkvJik= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2/go.mod h1:qwXFYgsP6T7XnJtbKlf1HP8AjxZZyzxMmc+Lq5GjlU4= github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4= github.com/datastax/go-cassandra-native-protocol v0.0.0-20211124104234-f6aea54fa801 h1:cxrfSL9aswWmLXsL6g1q02gln8GwH3sb4Fh2rGzKP8Q= github.com/datastax/go-cassandra-native-protocol v0.0.0-20211124104234-f6aea54fa801/go.mod h1:yFD0OKoVV9d1QW7Es58c1Gv6ijrqTGPcxgHv27wdC4Q= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/envoyproxy/go-control-plane v0.13.4 h1:zEqyPVyku6IvWCFwux4x9RxkLOMUL+1vC9xUFv5l2/M= -github.com/envoyproxy/go-control-plane v0.13.4/go.mod h1:kDfuBlDVsSj2MjrLEtRWtHlsWIFcGyB2RMO44Dc5GZA= -github.com/envoyproxy/go-control-plane/envoy v1.32.4 h1:jb83lalDRZSpPWW2Z7Mck/8kXZ5CQAFYVjQcdVIr83A= -github.com/envoyproxy/go-control-plane/envoy v1.32.4/go.mod h1:Gzjc5k8JcJswLjAx1Zm+wSYE20UrLtt7JZMWiWQXQEw= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/envoyproxy/go-control-plane v0.14.0 h1:hbG2kr4RuFj222B6+7T83thSPqLjwBIfQawTkC++2HA= +github.com/envoyproxy/go-control-plane v0.14.0/go.mod h1:NcS5X47pLl/hfqxU70yPwL9ZMkUlwlKxtAohpi2wBEU= +github.com/envoyproxy/go-control-plane/envoy v1.37.0 h1:u3riX6BoYRfF4Dr7dwSOroNfdSbEPe9Yyl09/B6wBrQ= +github.com/envoyproxy/go-control-plane/envoy v1.37.0/go.mod h1:DReE9MMrmecPy+YvQOAOHNYMALuowAnbjjEMkkWOi6A= github.com/envoyproxy/go-control-plane/ratelimit v0.1.0 h1:/G9QYbddjL25KvtKTv3an9lx6VBE2cnb8wp1vEGNYGI= github.com/envoyproxy/go-control-plane/ratelimit v0.1.0/go.mod h1:Wk+tMFAFbCXaJPzVVHnPgRKdUdwW/KdbRt94AzgRee4= -github.com/envoyproxy/protoc-gen-validate v1.2.1 h1:DEo3O99U8j4hBFwbJfrz9VtgcDfUKS7KJ7spH3d86P8= -github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU= +github.com/envoyproxy/protoc-gen-validate v1.3.3 h1:MVQghNeW+LZcmXe7SY1V36Z+WFMDjpqGAGacLe2T0ds= +github.com/envoyproxy/protoc-gen-validate v1.3.3/go.mod h1:TsndJ/ngyIdQRhMcVVGDDHINPLWB7C82oDArY51KfB0= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= +github.com/go-jose/go-jose/v4 v4.1.4 h1:moDMcTHmvE6Groj34emNPLs/qtYXRVcd6S7NHbHz3kA= +github.com/go-jose/go-jose/v4 v4.1.4/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= -github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= -github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/gocql/gocql v1.7.0 h1:O+7U7/1gSN7QTEAaMEsJc1Oq2QHXvCWoF3DFK9HDHus= @@ -70,12 +85,12 @@ github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4= -github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA= -github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q= -github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA= -github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys= -github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I= +github.com/googleapis/enterprise-certificate-proxy v0.3.16 h1:F/VPrx0YPBdksZJQdCAp0WUsqnNmZpUZszzfYt0M5Dw= +github.com/googleapis/enterprise-certificate-proxy v0.3.16/go.mod h1:9Yb0eAkH/Xqhvv3zbeKf/+wMJqCeocWc6KIhDvEAuYE= +github.com/googleapis/gax-go/v2 v2.22.0 h1:PjIWBpgGIVKGoCXuiCoP64altEJCj3/Ei+kSU5vlZD4= +github.com/googleapis/gax-go/v2 v2.22.0/go.mod h1:irWBbALSr0Sk3qlqb9SyJ1h68WjgeFuiOzI4Rqw5+aY= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 h1:8Tjv8EJ+pM1xP8mK6egEbD1OgnVTyacbefKhmbLhIhU= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2/go.mod h1:pkJQ2tZHJ0aFOVEEot6oZmaVEZcRme73eIFmhiVuRWs= github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed h1:5upAirOpQc1Q53c0bnx2ufif5kANL7bfZWcc6VJWJd8= github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed/go.mod h1:tMWxXQ9wFIaZeTI9F+hmhFiGpFmhOHzyShyFUhRm0H4= github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc= @@ -96,50 +111,51 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo= github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= -github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/rs/xid v1.2.1/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ= github.com/rs/zerolog v1.20.0/go.mod h1:IzD0RJ65iWH0w97OQQebJEvTZYvsCUm9WVLWBQrJRjo= +github.com/spiffe/go-spiffe/v2 v2.6.0 h1:l+DolpxNWYgruGQVV0xsfeya3CsC7m8iBzDnMpsbLuo= +github.com/spiffe/go-spiffe/v2 v2.6.0/go.mod h1:gm2SeUoMZEtpnzPNs2Csc0D/gX33k1xIx7lEzqblHEs= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= -github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= -github.com/tj/assert v0.0.3 h1:Df/BlaZ20mq6kuai7f5z2TvPFiwC3xaWJSDQNiIS3Rk= -github.com/tj/assert v0.0.3/go.mod h1:Ne6X72Q+TB1AteidzQncjw9PabbMp4PBMZ1k+vd1Pvk= -go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= -go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= -go.opentelemetry.io/contrib/detectors/gcp v1.34.0 h1:JRxssobiPg23otYU5SbWtQC//snGVIM3Tx6QRzlQBao= -go.opentelemetry.io/contrib/detectors/gcp v1.34.0/go.mod h1:cV4BMFcscUR/ckqLkbfQmF0PRsq8w/lMGzdbCSveBHo= -go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 h1:x7wzEgXfnzJcHDwStJT+mxOz4etr2EcexjqhBvmoakw= -go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0/go.mod h1:rg+RlpR5dKwaS95IyyZqj5Wd4E13lk/msnTS0Xl9lJM= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ= -go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ= -go.opentelemetry.io/otel v1.35.0/go.mod h1:UEqy8Zp11hpkUrL73gSlELM0DupHoiq72dR+Zqel/+Y= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/detectors/gcp v1.42.0 h1:kpt2PEJuOuqYkPcktfJqWWDjTEd/FNgrxcniL7kQrXQ= +go.opentelemetry.io/contrib/detectors/gcp v1.42.0/go.mod h1:W9zQ439utxymRrXsUOzZbFX4JhLxXU4+ZnCt8GG7yA8= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.68.0 h1:0Qx7VGBacMm9ZENQ7TnNObTYI4ShC+lHI16seduaxZo= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.68.0/go.mod h1:Sje3i3MjSPKTSPvVWCaL8ugBzJwik3u4smCjUeuupqg= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0 h1:CqXxU8VOmDefoh0+ztfGaymYbhdB/tT3zs79QaZTNGY= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0/go.mod h1:BuhAPThV8PBHBvg8ZzZ/Ok3idOdhWIodywz2xEcRbJo= +go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= +go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.24.0 h1:f2jriWfOdldanBwS9jNBdeOKAQN7b4ugAMaNu1/1k9g= go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.24.0/go.mod h1:B+bcQI1yTY+N0vqMpoZbEN7+XU4tNM0DmUiOwebFJWI= go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.31.0 h1:K0XaT3DwHAcV4nKLzcQvwAgSyisUghWoY20I7huthMk= go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.31.0/go.mod h1:B5Ki776z/MBnVha1Nzwp5arlzBbE3+1jk+pGmaP5HME= go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0 h1:Mw5xcxMwlqoJd97vwPxA8isEaIoxsta9/Q51+TTJLGE= go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0/go.mod h1:CQNu9bj7o7mC6U7+CA/schKEYakYXWr79ucDHTMGhCM= -go.opentelemetry.io/otel/metric v1.35.0 h1:0znxYu2SNyuMSQT4Y9WDWej0VpcsxkuklLa4/siN90M= -go.opentelemetry.io/otel/metric v1.35.0/go.mod h1:nKVFgxBZ2fReX6IlyW28MgZojkoAkJGaE8CpgeAU3oE= -go.opentelemetry.io/otel/sdk v1.35.0 h1:iPctf8iprVySXSKJffSS79eOjl9pvxV9ZqOWT0QejKY= -go.opentelemetry.io/otel/sdk v1.35.0/go.mod h1:+ga1bZliga3DxJ3CQGg3updiaAJoNECOgJREo9KHGQg= -go.opentelemetry.io/otel/sdk/metric v1.35.0 h1:1RriWBmCKgkeHEhM7a2uMjMUfP7MsOF5JpUCaEqEI9o= -go.opentelemetry.io/otel/sdk/metric v1.35.0/go.mod h1:is6XYCUMpcKi+ZsOvfluY5YstFnhW0BidkR+gL+qN+w= -go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt/xgMs= -go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc= -go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= -go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8= -go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= -go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= +go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= +go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= +go.opentelemetry.io/otel/metric/x v0.66.0 h1:YkCrx1zLOChi9ZcZ6euupOcsgzbVlec7D/xoEU1+cTA= +go.opentelemetry.io/otel/metric/x v0.66.0/go.mod h1:d1+BDj9t96do0/1LoU1ayfCv79ZgNE41qbhBvnMOBZk= +go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58= +go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0= +go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI= +go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA= +go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= +go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= +go.opentelemetry.io/proto/otlp v1.9.0 h1:l706jCMITVouPOqEnii2fIAuO3IVGBRPV5ICjceRb/A= +go.opentelemetry.io/proto/otlp v1.9.0/go.mod h1:xE+Cx5E/eEHw+ISFkwPLwCZefwVjY+pqKg1qcK03+/4= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= @@ -147,40 +163,42 @@ go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN8 go.uber.org/zap v1.26.0 h1:sI7k6L95XOKS281NhVKOFCUNIvv9e0w4BF8N3u+tCRo= go.uber.org/zap v1.26.0/go.mod h1:dtElttAiwGvoJ/vj4IwHBS/gXsEu/pZ50mUIRWuG0so= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= -golang.org/x/crypto v0.36.0 h1:AnAEvhDddvBdpY+uR+MyHmuZzzNqXSe/GvuDeob5L34= -golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc= +golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI= +golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8= golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63 h1:m64FZMko/V45gv0bNmrNYoDEq8U5YUhetc9cBWKS1TQ= golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63/go.mod h1:0v4NqG35kSWCMzLaMeX+IQrlSnVE/bqGSyC2cz/9Le8= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= -golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8= -golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= -golang.org/x/oauth2 v0.28.0 h1:CrgCKl8PPAVtLnU3c+EDw6x11699EWlsDeWNWKdIOkc= -golang.org/x/oauth2 v0.28.0/go.mod h1:onh5ek6nERTohokkhCD/y2cV4Do3fxFHFuAejCkRWT8= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.12.0 h1:MHc5BpPuC30uJk597Ri8TV3CNZcTLu6B6z4lJy+g6Jw= -golang.org/x/sync v0.12.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= -golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= -golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= -golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY= -golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4= -golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0= -golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.0.0-20190828213141-aed303cbaa74/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -google.golang.org/api v0.228.0 h1:X2DJ/uoWGnY5obVjewbp8icSL5U4FzuCfy9OjbLSnLs= -google.golang.org/api v0.228.0/go.mod h1:wNvRS1Pbe8r4+IfBIniV8fwCpGwTrYa+kMUDiC5z5a4= -google.golang.org/genproto v0.0.0-20250404141209-ee84b53bf3d0 h1:wX+y2uwLyC73sX9zfiJW7E7m68+oxAQGzgCmoM0e/zs= -google.golang.org/genproto v0.0.0-20250404141209-ee84b53bf3d0/go.mod h1:jwIveCnYVWLDIe0ZXnIrfMKNoy/rQRSRrepUPEruz0U= -google.golang.org/genproto/googleapis/api v0.0.0-20250404141209-ee84b53bf3d0 h1:Qbb5RVn5xzI4naMJSpJ7lhvmos6UwZkbekd5Uz7rt9E= -google.golang.org/genproto/googleapis/api v0.0.0-20250404141209-ee84b53bf3d0/go.mod h1:6T35kB3IPpdw7Wul09by0G/JuOuIFkXV6OOvt8IZeT8= -google.golang.org/genproto/googleapis/rpc v0.0.0-20250404141209-ee84b53bf3d0 h1:0K7wTWyzxZ7J+L47+LbFogJW1nn/gnnMCN0vGXNYtTI= -google.golang.org/genproto/googleapis/rpc v0.0.0-20250404141209-ee84b53bf3d0/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A= -google.golang.org/grpc v1.71.1 h1:ffsFWr7ygTUscGPI0KKK6TLrGz0476KUvvsbqWK0rPI= -google.golang.org/grpc v1.71.1/go.mod h1:H0GRtasmQOh9LkFoCPDu3ZrwUtD1YGE+b2vYBYd/8Ec= -google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY= -google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/api v0.283.0 h1:0lkp8u0MPwJVHqRL+nJlMAoZVVzbmiXmFHXMOTmSPik= +google.golang.org/api v0.283.0/go.mod h1:6Wssta4c5n9qHq5CBhmlai5h/PUa1djdDAIhYEHyvcM= +google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94 h1:YJjbgu+dkp5kUJLfpMyCLfBIWZb/FcJyuLeo1gVBOuo= +google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94/go.mod h1:RRHjglSYABVCWpQ7USCpdfhcd9t4PkajvVwyynZizTc= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1:Kjn0N0tCrDgiAFW+lGO4JZ3ck44CehvJQMAwj9QF0G8= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:q4lMZS6kskjT5HvCPrnnypcDPVJqT/f4nfxmkE7gryY= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= +google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= @@ -192,7 +210,6 @@ gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYs gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -gopkg.in/yaml.v3 v3.0.0-20200605160147-a5ece683394c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/mem_table/engine_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/mem_table/engine_test.go index e93d6fdc..b615692e 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/mem_table/engine_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/mem_table/engine_test.go @@ -8,6 +8,7 @@ import ( "github.com/datastax/go-cassandra-native-protocol/primitive" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/zap" "log" "testing" "time" @@ -200,8 +201,8 @@ func Test_SelectEngine(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - tr := select_translator.NewSelectTranslator(schemas) - preparedQuery, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeSelect), tt.sessionKeyspace) + tr := select_translator.NewSelectTranslator(schemas, zap.NewNop()) + preparedQuery, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeSelect), tt.sessionKeyspace) require.NoError(t, err) values := types.NewQueryParameterValues(preparedQuery.Parameters(), time.Now()) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata/utils.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata/utils.go index 3efb34aa..64a8dca7 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata/utils.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata/utils.go @@ -87,7 +87,8 @@ func createBigtableRowKeySchema(primaryKeys []types.CreateTablePrimaryKeyConfig, rowKeySchemaFields = append(rowKeySchemaFields, part) } return &bigtable.StructType{ - Fields: rowKeySchemaFields, + Fields: rowKeySchemaFields, + Encoding: bigtable.StructOrderedCodeBytesEncoding{}, }, nil } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/README.md b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/README.md deleted file mode 100644 index 302f9528..00000000 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/README.md +++ /dev/null @@ -1,159 +0,0 @@ - -# Integrating OpenTelemetry (OTEL) with Your Application - -## Overview - -OpenTelemetry (OTEL) provides a set of tools, APIs, and SDKs to instrument, generate, collect, and export telemetry data (metrics, logs, and traces) for monitoring and troubleshooting your applications. This guide will help you integrate OTEL into your application, enabling you to capture and export metrics and traces. - - -# Setting up the OTEL Collector Service on GCP - -This guide provides step-by-step instructions to set up the OpenTelemetry (OTEL) Collector Service on Google Cloud Platform (GCP) using the provided configuration file. - -## Prerequisites - -Before enabling OTEL in your application, ensure that the collector service is up and running. The collector service is responsible for capturing and exporting the telemetry data. By default, OTEL is disabled in the configuration. - -1. **GCP Account**: Ensure you have a GCP account and have access to a GCP project. -2. **Google Cloud SDK**: Install the [Google Cloud SDK](https://cloud.google.com/sdk/docs/install). -3. **Docker**: Install Docker to run the OTEL Collector in a container. -4. **Kubernetes Cluster (GKE)**: Set up a Google Kubernetes Engine (GKE) cluster if deploying in a Kubernetes environment. - -## Steps to Set Up and Configure the OTEL Collector Service with Proxy Adaptor. - -### Step 1: Use the below config.yaml (or set up the `CONFIG_FILE` environment variable) file for OTEL collector service - -```yaml - # Enable the health check in the proxy application config only if the - # "health_check" extension is added to this OTEL config for the collector service. - # - # Recommendation: - # Enable the OTEL health check if you need to verify the collector's availability - # at the start of the application. For development or testing environments, it can - # be safely disabled to reduce complexity. -receivers: - otlp: - protocols: - grpc: - endpoint: "0.0.0.0:4317" - http: - endpoint: "0.0.0.0:55681" - -processors: - batch: - send_batch_max_size: 0 - send_batch_size: 8192 - timeout: 5s - - memory_limiter: - # drop metrics if memory usage gets too high - check_interval: 5s - limit_percentage: 65 - spike_limit_percentage: 20 - - resourcedetection: - detectors: [gcp] - timeout: 10s - -exporters: - googlecloud: - metric: - instrumentation_library_labels: true - service_resource_labels: true - -extensions: - health_check: - endpoint: "0.0.0.0:13133" - -service: - extensions: [health_check] - pipelines: - metrics: - receivers: [otlp] - processors: [batch, memory_limiter, resourcedetection] - exporters: [googlecloud] - traces: - receivers: [otlp] - processors: [batch, memory_limiter, resourcedetection] - exporters: [googlecloud] -``` - -### Step 2: Configure Proxy Adaptor - -Follow these steps to enable and configure OTEL in your application: - -1. **Edit the `config.yaml` (or set up the `CONFIG_FILE` environment variable) File:** - - Set the `enabled` field to `True` to enable metrics and traces. - - Configure the endpoints for the collector service to export the metrics and traces. - - Set the healthcheck endpoint, which is configured in the collector service. Avoid including `http://` in the endpoints. Refer to the `example_config.yaml` file for guidance. - -2. **Example Configuration Block for OTEL:** - ```yaml - # Enable the health check in this proxy application config only if the - # "health_check" extension is added to the OTEL collector service configuration. - # - # Recommendation: - # Enable the OTEL health check if you need to verify the collector's availability - # at the start of the application. For development or testing environments, it can - # be safely disabled to reduce complexity. - otel: - # Set enabled to true or false for OTEL metrics and traces - enabled: True - # Name of the collector service to be setup as a sidecar - serviceName: YOUR_OTEL_COLLECTOR_SERVICE_NAME - healthcheck: - # Enable/Disable Health Check for OTEL, Default 'False'. - enabled: False - # Health check endpoint for the OTEL collector service - endpoint: YOUR_OTEL_COLLECTOR_HEALTHCHECK_ENDPOINT - metrics: - # Collector service endpoint - endpoint: YOUR_OTEL_COLLECTOR_SERVICE_ENDPOINT - traces: - # Collector service endpoint - endpoint: YOUR_OTEL_COLLECTOR_SERVICE_ENDPOINT - # Sampling ratio should be between 0 and 1. - samplingRatio: YOUR_SAMPLING_RATIO - - -### Step 3: Setup Proxy Adaptor & OTEL Collector service as a sidecar on GKE. - -1. Use the `proxy-adapter-application-as-sidecar.yaml` file from the deployment/sidecar-k8 to setup the otel collector service along with proxy adaptor as a sidecar. You can find step by step instructions in the `/deployment/sidecar-k8/README.md` file. - - -### Step 4: Verify OTEL changes - -1. Verify the Health Check: - -Run a curl command from the k8 pod to http://collector-ip:13133/health_check. You should see a health status message. - -2. Check OTEL traces on GCP - -- Login to your GCP project and open Monitoring service. -- Open the sidebar on the left and click on Trace Explorer -- You should be able to see the traces as shown below: - -![Alt text](./img/traces-execute.png) -![Alt text](./img/traces-batch.png) - -3. Check OTEL metrics on GCP - -- Login to your GCP project and open Monitoring service. -- Open the sidebar on the left and click on Metrics Explorer -- You should be able to see the categories as shown below: - -![Alt text](./img/metrics-category.png) - -- Select options `Prometheus Target/Bigtable` -- To view the metrics for total number of requests, select option -`prometheus/bigtable_cassandra_adapter_request_count_total/counter` - -![Alt text](./img/metrics_total_requests.png) - -- To view the metrics for latency, select option -`prometheus/bigtable_cassandra_adapter_roundtrip_latencies_milliseconds/histogram` - -![Alt text](./img/metrics-latency.png) - -- You can also view other metrics related to Bigtable library under the `Prometheus Target/Bigtable` category. - diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/metrics.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/metrics.go new file mode 100644 index 00000000..f9fa1ba5 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/metrics.go @@ -0,0 +1,84 @@ +/* + * Copyright (C) 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy of + * the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ + +package otelgo + +import ( + "context" + "errors" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "go.opentelemetry.io/otel/metric/noop" + "time" + + metricexporter "github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc" + "go.opentelemetry.io/otel/metric" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/resource" +) + +const ( + requestCountMetric = "bigtable/cassandra_adapter/request_count" + latencyMetric = "bigtable/cassandra_adapter/roundtrip_latencies" +) + +func InitMeterProvider(ctx context.Context, config *types.OtelConfig, res *resource.Resource) (metric.MeterProvider, ShutdownFn, error) { + if !config.Enabled { + return noop.NewMeterProvider(), func(_ context.Context) error { return nil }, nil + } + + var exporter sdkmetric.Exporter + var err error + + if config.Metrics.GcpMetricsEnabled { + exporter, err = metricexporter.New(metricexporter.WithProjectID(config.Traces.ProjectId)) + } else { + if !isValidEndpoint(config.Metrics.Endpoint) { + return nil, nil, errors.New("invalid metric endpoint format") + } + exporter, err = otlpmetricgrpc.New(ctx, otlpmetricgrpc.WithEndpoint(config.Metrics.Endpoint), otlpmetricgrpc.WithInsecure()) + } + + if err != nil { + return nil, nil, err + } + mp := sdkmetric.NewMeterProvider( + sdkmetric.WithReader(sdkmetric.NewPeriodicReader(exporter)), + sdkmetric.WithResource(res), + sdkmetric.WithView(sdkmetric.NewView( + sdkmetric.Instrument{Name: "rpc.client.*"}, + sdkmetric.Stream{Aggregation: sdkmetric.AggregationDrop{}}, + )), + ) + return mp, mp.Shutdown, nil +} + +func (o *OpenTelemetry) RecordMetrics(ctx context.Context, startTime time.Time, attrs types.Attributes) { + a := CommonAttributes(attrs) + o.requestLatency.Record(ctx, time.Since(startTime).Milliseconds(), metric.WithAttributes(a...)) + o.requestCount.Add(ctx, 1, metric.WithAttributes(a...)) +} + +func CommonAttributes(attrs types.Attributes) []attribute.KeyValue { + return []attribute.KeyValue{ + attributeKeyKeyspace.String(string(attrs.Keyspace)), + attributeKeyMethod.String(attrs.Method), + attributeKeyQueryType.String(attrs.QueryType.String()), + attributeKeyTable.String(string(attrs.Table)), + attributeKeyStatus.String(attrs.Status), + } +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel.go index a3722d44..6b67d391 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel.go @@ -18,132 +18,73 @@ package otelgo import ( "context" - "errors" + "fmt" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "go.opentelemetry.io/otel/trace" "net/http" "net/url" "strings" - "time" "github.com/google/uuid" "go.opentelemetry.io/contrib/detectors/gcp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/codes" - "go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc" - "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" "go.opentelemetry.io/otel/metric" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" "go.opentelemetry.io/otel/sdk/resource" - sdktrace "go.opentelemetry.io/otel/sdk/trace" semconv "go.opentelemetry.io/otel/semconv/v1.4.0" - "go.opentelemetry.io/otel/trace" "go.uber.org/zap" ) -type Attributes struct { - Method string - Status string - QueryType string - Keyspace string -} +type ShutdownFn func(ctx context.Context) error var ( - attributeKeyDatabase = attribute.Key("database") attributeKeyMethod = attribute.Key("method") attributeKeyStatus = attribute.Key("status") - attributeKeyInstance = attribute.Key("instance") + attributeKeyKeyspace = attribute.Key("keyspace") attributeKeyQueryType = attribute.Key("query_type") + attributeKeyTable = attribute.Key("table") ) -// OTelConfig holds configuration for OpenTelemetry. -type OTelConfig struct { - TracerEndpoint string - MetricEndpoint string - ServiceName string - TraceSampleRatio float64 - OTELEnabled bool - Database string - Instance string - HealthCheckEnabled bool - HealthCheckEp string - ServiceVersion string -} - -const ( - requestCountMetric = "bigtable/cassandra_adapter/request_count" - latencyMetric = "bigtable/cassandra_adapter/roundtrip_latencies" -) - -// OpenTelemetry provides methods to setup tracing and metrics. type OpenTelemetry struct { - Config *OTelConfig - tracer trace.Tracer + Config *types.OtelConfig requestCount metric.Int64Counter requestLatency metric.Int64Histogram logger *zap.Logger } -// NewOpenTelemetry() initializes OpenTelemetry tracing and metrics components. -// It sets up the tracer and meter providers, configures health checks (if enabled), -// and returns an OpenTelemetry instance along with a shutdown function. -// -// Parameters: -// - ctx: Context for managing OpenTelemetry lifecycle. -// - config: Configuration struct for OpenTelemetry settings. -// - logger: Logger instance for capturing OpenTelemetry logs. -// -// Returns: -// - *OpenTelemetry: A configured instance of OpenTelemetry. -// - func(context.Context) error: A shutdown function to clean up resources. -// - error: An error if initialization fails. -func NewOpenTelemetry(ctx context.Context, config *OTelConfig, logger *zap.Logger) (*OpenTelemetry, func(context.Context) error, error) { +// NewOpenTelemetry initializes OpenTelemetry tracing and metrics components. +func NewOpenTelemetry(ctx context.Context, config *types.OtelConfig, logger *zap.Logger) (*OpenTelemetry, ShutdownFn, error) { otelInst := &OpenTelemetry{Config: config, logger: logger} - var err error - otelInst.Config.OTELEnabled = config.OTELEnabled - if !config.OTELEnabled { - return otelInst, nil, nil - } - if config.HealthCheckEnabled { - resp, err := http.Get("http://" + config.HealthCheckEp) - if err != nil { - return nil, nil, err - } - if resp.StatusCode != 200 { - return nil, nil, errors.New("OTEL collector service is not up and running") + if config.Enabled && config.HealthCheck.Enabled { + resp, err := http.Get("http://" + config.HealthCheck.Endpoint) + if err != nil || resp.StatusCode != 200 { + return nil, nil, fmt.Errorf("OTEL health check failed: %v", err) } logger.Info("OTEL health check complete") } - var shutdownFuncs []func(context.Context) error - otelResource := buildOtelResource(ctx, config) - // Initialize tracerProvider - tracerProvider, err := InitTracerProvider(ctx, config, otelResource) + res := buildOtelResource(ctx, config) + + tp, shutdownTp, err := createTraceProvider(ctx, config, res) if err != nil { - logger.Error("error while initializing the tracer provider", zap.Error(err)) - return nil, nil, err + return nil, nil, fmt.Errorf("failed to create trace provider: %w", err) } - otel.SetTracerProvider(tracerProvider) - otelInst.tracer = tracerProvider.Tracer(config.ServiceName) - shutdownFuncs = append(shutdownFuncs, tracerProvider.Shutdown) + otel.SetTracerProvider(tp) - // Initialize MeterProvider - meterProvider, err := InitMeterProvider(ctx, config, otelResource) + mp, shutdownMp, err := InitMeterProvider(ctx, config, res) if err != nil { - logger.Error("error while initializing the meter provider", zap.Error(err)) - return nil, nil, err + return nil, nil, fmt.Errorf("failed to create meter provider: %w", err) } - otel.SetMeterProvider(meterProvider) - meter := meterProvider.Meter(config.ServiceName) - shutdownFuncs = append(shutdownFuncs, meterProvider.Shutdown) - shutdown := shutdownOpenTelemetryComponents(shutdownFuncs) - otelInst.requestCount, err = meter.Int64Counter(requestCountMetric, metric.WithDescription("Records metric for number of query requests coming in"), metric.WithUnit("1")) - if err != nil { - logger.Error("error during registering instrument for metric bigtable/cassandra_adapter/request_count", zap.Error(err)) - return nil, nil, err + otel.SetMeterProvider(mp) + meter := mp.Meter(config.ServiceName) + + var metricsErr error + otelInst.requestCount, metricsErr = meter.Int64Counter(requestCountMetric, metric.WithDescription("Records number of query requests"), metric.WithUnit("1")) + if metricsErr != nil { + return nil, nil, metricsErr } - otelInst.requestLatency, err = meter.Int64Histogram(latencyMetric, + otelInst.requestLatency, metricsErr = meter.Int64Histogram(latencyMetric, metric.WithDescription("Records latency for all query operations"), metric.WithExplicitBucketBoundaries(0.0, 0.0010, 0.0013, 0.0016, 0.0020, 0.0024, 0.0031, 0.0038, 0.0048, 0.0060, 0.0075, 0.0093, 0.0116, 0.0146, 0.0182, 0.0227, 0.0284, 0.0355, 0.0444, 0.0555, 0.0694, 0.0867, @@ -152,316 +93,56 @@ func NewOpenTelemetry(ctx context.Context, config *OTelConfig, logger *zap.Logge 22.9589, 28.6986, 35.8732, 44.8416, 56.0519, 70.0649, 87.5812, 109.4764, 136.8456, 171.0569, 213.8212, 267.2765, 334.0956, 417.6195, 522.0244, 652.5304), metric.WithUnit("ms")) - if err != nil { - logger.Error("error during registering instrument for metric bigtable/cassandra_adapter/roundtrip_latencies", zap.Error(err)) - return nil, nil, err + if metricsErr != nil { + return nil, nil, metricsErr } - return otelInst, shutdown, nil -} -// shutdownOpenTelemetryComponents() aggregates multiple shutdown functions into a single callable function. -// It iterates over all shutdown functions, executing them sequentially. -// -// Parameters: -// - shutdownFuncs: A slice of shutdown functions for OpenTelemetry components. -// -// Returns: -// - func(context.Context) error: A single shutdown function that cleans up all initialized components. -func shutdownOpenTelemetryComponents(shutdownFuncs []func(context.Context) error) func(context.Context) error { - return func(ctx context.Context) error { - var shutdownErr error - for _, shutdownFunc := range shutdownFuncs { - if err := shutdownFunc(ctx); err != nil { - shutdownErr = err - } + shutdown := func(ctx context.Context) error { + err1 := shutdownTp(ctx) + err2 := shutdownMp(ctx) + if err1 != nil { + return err1 } - return shutdownErr + return err2 } -} -// InitTracerProvider() configures and initializes an OpenTelemetry TracerProvider. -// It sets up a gRPC-based OTLP trace exporter and applies the sampling strategy. -// -// Parameters: -// - ctx: Context for managing initialization. -// - config: OpenTelemetry configuration settings. -// - resource: OpenTelemetry resource with metadata. -// -// Returns: -// - *sdktrace.TracerProvider: A configured TracerProvider instance. -// - error: An error if initialization fails. -func InitTracerProvider(ctx context.Context, config *OTelConfig, resource *resource.Resource) (*sdktrace.TracerProvider, error) { - sampler := sdktrace.TraceIDRatioBased(config.TraceSampleRatio) - if config.TracerEndpoint == "" { - return nil, errors.New("tracer endpoint cannot be empty") - } - // Basic validation for incorrect endpoint format - if !isValidEndpoint(config.TracerEndpoint) { - return nil, errors.New("invalid tracer endpoint format") - } - - traceExporter, err := otlptracegrpc.New(ctx, - otlptracegrpc.WithEndpoint(config.TracerEndpoint), - otlptracegrpc.WithInsecure(), - ) - if err != nil { - return nil, err - } - - tp := sdktrace.NewTracerProvider( - sdktrace.WithBatcher(traceExporter), - sdktrace.WithResource(resource), - sdktrace.WithSampler(sdktrace.ParentBased(sampler)), - ) - return tp, nil + return otelInst, shutdown, nil } -// InitMeterProvider() initializes an OpenTelemetry MeterProvider for collecting application metrics. -// It configures a gRPC exporter to send metrics data and applies filtering to exclude unnecessary gRPC metrics. -// -// Parameters: -// - ctx: Context for managing initialization. -// - config: OpenTelemetry configuration settings. -// - resource: OpenTelemetry resource with metadata. -// -// Returns: -// - *sdkmetric.MeterProvider: A configured MeterProvider instance. -// - error: An error if initialization fails. -func InitMeterProvider(ctx context.Context, config *OTelConfig, resource *resource.Resource) (*sdkmetric.MeterProvider, error) { - if config.MetricEndpoint == "" { - return nil, errors.New("metric endpoint cannot be empty") - } - - // Basic validation for incorrect endpoint format - if !isValidEndpoint(config.MetricEndpoint) { - return nil, errors.New("invalid tracer endpoint format") - } - var views []sdkmetric.View - me, err := otlpmetricgrpc.New(ctx, - otlpmetricgrpc.WithEndpoint(config.MetricEndpoint), - otlpmetricgrpc.WithInsecure(), - ) - if err != nil { - return nil, err +func buildOtelResource(ctx context.Context, config *types.OtelConfig) *resource.Resource { + attrs := []attribute.KeyValue{ + semconv.ServiceNameKey.String(config.ServiceName), + semconv.ServiceInstanceIDKey.String(uuid.New().String()), + semconv.ServiceVersionKey.String(config.ServiceVersion), } - - // Define views to filter out unwanted gRPC metrics - views = []sdkmetric.View{ - sdkmetric.NewView( - sdkmetric.Instrument{Name: "rpc.client.*"}, // Wildcard pattern to match gRPC client metrics - sdkmetric.Stream{Aggregation: sdkmetric.AggregationDrop{}}, // Drop these metrics - )} - - mp := sdkmetric.NewMeterProvider( - sdkmetric.WithReader(sdkmetric.NewPeriodicReader(me)), - sdkmetric.WithResource(resource), - sdkmetric.WithView(views...), - ) - return mp, nil -} - -// buildOtelResource() creates an OpenTelemetry resource containing metadata about the service. -// It uses GCP resource detectors and falls back to manually provided attributes if necessary. -// -// Parameters: -// - ctx: Context for managing initialization. -// - config: OpenTelemetry configuration settings. -// -// Returns: -// - *resource.Resource: A configured OpenTelemetry resource containing metadata. -func buildOtelResource(ctx context.Context, config *OTelConfig) *resource.Resource { res, err := resource.New(ctx, resource.WithSchemaURL(semconv.SchemaURL), - // Use the GCP resource detector! resource.WithDetectors(gcp.NewDetector()), - // Keep the default detectors resource.WithTelemetrySDK(), - resource.WithAttributes( - semconv.ServiceNameKey.String(config.ServiceName), - semconv.ServiceInstanceIDKey.String(uuid.New().String()), - semconv.ServiceVersionKey.String(config.ServiceVersion), - ), + resource.WithAttributes(attrs...), ) - if err != nil { - // Default resource - return resource.NewWithAttributes( - semconv.SchemaURL, - semconv.ServiceNameKey.String(config.ServiceName), - semconv.ServiceInstanceIDKey.String(uuid.New().String()), - semconv.ServiceVersionKey.String(config.ServiceVersion), - ) + return resource.NewWithAttributes(semconv.SchemaURL, attrs...) } - return res - } -// StartSpan() creates and starts a new trace span in OpenTelemetry. -// If OpenTelemetry is disabled, it returns the original context. -// -// Parameters: -// - ctx: The current execution context. -// - name: The name of the span to be created. -// - attrs: A list of attributes to associate with the span. -// -// Returns: -// - context.Context: The updated context containing the new span. -// - trace.Span: The created span instance. -func (o *OpenTelemetry) StartSpan(ctx context.Context, name string, attrs []attribute.KeyValue) (context.Context, trace.Span) { - if !o.Config.OTELEnabled { - return ctx, nil - } - - ctx, span := o.tracer.Start(ctx, name, trace.WithAttributes(attrs...)) - return ctx, span -} - -// RecordError() logs an error inside an active trace span in OpenTelemetry. -// It updates the span's status to indicate an error has occurred. -// -// Parameters: -// - span: The active trace span where the error should be recorded. -// - err: The error to be recorded in the span. If nil, the span is marked as OK. -func (o *OpenTelemetry) RecordError(span trace.Span, err error) { - if !o.Config.OTELEnabled { - return - } - if err != nil { - span.RecordError(err) - span.SetStatus(codes.Error, err.Error()) - } else { - span.SetStatus(codes.Ok, "") - } -} - -// EndSpan() finalizes the current span in OpenTelemetry. -// If OpenTelemetry is disabled, this function does nothing. -// -// Parameters: -// - span: The span to be ended. -func (o *OpenTelemetry) EndSpan(span trace.Span) { - if !o.Config.OTELEnabled { - return - } - - span.End() -} - -// RecordMetrics() records request count and latency metrics in OpenTelemetry. -// It determines whether the request was successful or failed based on the error parameter. -// -// Parameters: -// - ctx: The execution context for OpenTelemetry. -// - method: The name of the method being recorded. -// - startTime: The start time of the request for latency calculation. -// - queryType: The type of query being executed (e.g., "select", "insert"). -// - err: The error encountered, if any. Used to determine success/failure status. -func (o *OpenTelemetry) RecordMetrics(ctx context.Context, method string, startTime time.Time, queryType string, keyspace types.Keyspace, err error) { - status := "OK" - if err != nil { - status = "failure" - } - o.RecordRequestCountMetric(ctx, Attributes{ - Method: method, - Status: status, - QueryType: queryType, - Keyspace: string(keyspace), - }) - o.RecordLatencyMetric(ctx, startTime, Attributes{ - Method: method, - QueryType: queryType, - Keyspace: string(keyspace), - }) -} - -// RecordLatencyMetric() records the latency of an operation in OpenTelemetry. -// It dynamically builds metric attributes before sending the recorded value. -// -// Parameters: -// - ctx: The execution context. -// - startTime: The time when the operation started, used for latency calculation. -// - attrs: Additional attributes to associate with the latency metric. -func (o *OpenTelemetry) RecordLatencyMetric(ctx context.Context, startTime time.Time, attrs Attributes) { - if !o.Config.OTELEnabled { - return - } - - // Build attributes dynamically - attr := []attribute.KeyValue{ - attributeKeyInstance.String(attrs.Keyspace), - attributeKeyDatabase.String(o.Config.Database), - attributeKeyMethod.String(attrs.Method), - attributeKeyQueryType.String(attrs.QueryType), - } - attr = append(attr, attributeKeyMethod.String(attrs.Method)) - attr = append(attr, attributeKeyQueryType.String(attrs.QueryType)) - o.requestLatency.Record(ctx, int64(time.Since(startTime).Milliseconds()), metric.WithAttributes(attr...)) -} - -// RecordRequestCountMetric() increments the request count metric in OpenTelemetry. -// It dynamically builds metric attributes before sending the recorded value. -// -// Parameters: -// - ctx: The execution context. -// - attrs: Attributes associated with the request (e.g., method, status). -func (o *OpenTelemetry) RecordRequestCountMetric(ctx context.Context, attrs Attributes) { - if !o.Config.OTELEnabled { - return - } - - // Build attributes dynamically - attr := []attribute.KeyValue{ - attributeKeyInstance.String(attrs.Keyspace), - attributeKeyDatabase.String(o.Config.Database), - attributeKeyMethod.String(attrs.Method), - attributeKeyQueryType.String(attrs.QueryType), - attributeKeyStatus.String(attrs.Status), - } - attr = append(attr, attributeKeyMethod.String(attrs.Method)) - attr = append(attr, attributeKeyQueryType.String(attrs.QueryType)) - attr = append(attr, attributeKeyStatus.String(attrs.Status)) - o.requestCount.Add(ctx, 1, metric.WithAttributes(attr...)) -} - -// AddAnnotation() adds an event annotation to the active span in the given context. -// -// Parameters: -// - ctx: The execution context containing the span. -// - event: The event name to be added as an annotation. -func AddAnnotation(ctx context.Context, event string) { - span := trace.SpanFromContext(ctx) - span.AddEvent(event) -} - -// AddAnnotationWithAttr() adds an event annotation with attributes to the active span. -// -// Parameters: -// - ctx: The execution context containing the span. -// - event: The event name to be added as an annotation. -// - attr: A list of attributes to attach to the annotation. -func AddAnnotationWithAttr(ctx context.Context, event string, attr []attribute.KeyValue) { - span := trace.SpanFromContext(ctx) - span.AddEvent(event, trace.WithAttributes(attr...)) -} - -// isValidEndpoint checks if the given endpoint is a valid host:port format func isValidEndpoint(endpoint string) bool { + if endpoint == "" { + return false + } if strings.Contains(endpoint, "://") { - parsedURL, err := url.Parse(endpoint) - if err != nil { - return false - } - // Check if the original endpoint string had an empty host. - if strings.HasPrefix(endpoint, parsedURL.Scheme+"://:") { - return false - } - if parsedURL.Host == "" || parsedURL.Port() == "" { - return false - } - return true + u, err := url.Parse(endpoint) + return err == nil && u.Host != "" && u.Port() != "" } - parts := strings.Split(endpoint, ":") return len(parts) == 2 && parts[0] != "" && parts[1] != "" } + +func AddQueryAnnotations(s trace.Span, q types.IQuery) { + s.SetAttributes( + attributeKeyQueryType.String(q.QueryType().String()), + attributeKeyKeyspace.String(string(q.Keyspace())), + attributeKeyTable.String(string(q.Table())), + ) +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel_test.go deleted file mode 100644 index e7a907e1..00000000 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/otel_test.go +++ /dev/null @@ -1,971 +0,0 @@ -/* - * Copyright (C) 2025 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); you may not - * use this file except in compliance with the License. You may obtain a copy of - * the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT - * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the - * License for the specific language governing permissions and limitations under - * the License. - */ - -package otelgo - -import ( - "context" - "errors" - "fmt" - "net/http" - "reflect" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "go.opentelemetry.io/otel/attribute" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" - "go.opentelemetry.io/otel/sdk/resource" - sdktrace "go.opentelemetry.io/otel/sdk/trace" - semconv "go.opentelemetry.io/otel/semconv/v1.4.0" - "go.opentelemetry.io/otel/trace" - "go.opentelemetry.io/otel/trace/noop" - "go.uber.org/zap" -) - -func TestNewOpenTelemetry(t *testing.T) { - ctx := context.Background() - srv1 := setupTestEndpoint(":7060", "/trace") - srv2 := setupTestEndpoint(":7061", "/metric") - srv := setupTestEndpoint(":7062", "/TestNewOpenTelemetry") - time.Sleep(3 * time.Second) - defer func() { - assert.NoError(t, srv.Shutdown(ctx), "failed to shutdown srv") - assert.NoError(t, srv1.Shutdown(ctx), "failed to shutdown srv1") - assert.NoError(t, srv2.Shutdown(ctx), "failed to shutdown srv2") - }() - type args struct { - ctx context.Context - config *OTelConfig - logger *zap.Logger - } - tests := []struct { - name string - args args - wantErr bool - }{ - { - name: "Test success", - args: args{ - ctx: ctx, - config: &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7062/TestNewOpenTelemetry", - }, - logger: zap.NewNop(), - }, - wantErr: false, - }, - { - name: "Test when otel disabled", - args: args{ - ctx: ctx, - config: &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: false, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7062/TestNewOpenTelemetry", - }, - logger: zap.NewNop(), - }, - wantErr: false, - }, - { - name: "Test when healthcheck endpoint missing", - args: args{ - ctx: ctx, - config: &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEnabled: false, - HealthCheckEp: "", - }, - logger: zap.NewNop(), - }, - wantErr: false, - }, - { - name: "Test error when healthcheck endpoint missing", - args: args{ - ctx: ctx, - config: &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEnabled: true, - HealthCheckEp: "", - }, - logger: zap.NewNop(), - }, - wantErr: true, - }, - { - name: "Test error when TracerEndpoint endpoint missing", - args: args{ - ctx: ctx, - config: &OTelConfig{ - TracerEndpoint: "", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7062/TestNewOpenTelemetry", - }, - logger: zap.NewNop(), - }, - wantErr: true, - }, - { - name: "Test when MetricEndpoint endpoint missing", - args: args{ - ctx: ctx, - config: &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7062/TestNewOpenTelemetry", - }, - logger: zap.NewNop(), - }, - wantErr: true, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - - _, _, err := NewOpenTelemetry(tt.args.ctx, tt.args.config, tt.args.logger) - if (err != nil) != tt.wantErr { - t.Errorf("NewOpenTelemetry() error = %v, wantErr %v", err, tt.wantErr) - return - } - }) - } -} - -func TestShutdownOpenTelemetryComponents(t *testing.T) { - t.Run("All functions succeed", func(t *testing.T) { - ctx := context.Background() - - shutdownFunc1 := func(ctx context.Context) error { - return nil - } - shutdownFunc2 := func(ctx context.Context) error { - return nil - } - - shutdownFuncs := []func(context.Context) error{ - shutdownFunc1, - shutdownFunc2, - } - - shutdown := shutdownOpenTelemetryComponents(shutdownFuncs) - err := shutdown(ctx) - assert.NoError(t, err) - }) - - t.Run("One function fails", func(t *testing.T) { - ctx := context.Background() - - shutdownFunc1 := func(ctx context.Context) error { - return nil - } - shutdownFunc2 := func(ctx context.Context) error { - return errors.New("shutdown error") - } - - shutdownFuncs := []func(context.Context) error{ - shutdownFunc1, - shutdownFunc2, - } - - shutdown := shutdownOpenTelemetryComponents(shutdownFuncs) - err := shutdown(ctx) - assert.Error(t, err) - assert.Equal(t, "shutdown error", err.Error()) - }) - - t.Run("Multiple functions fail", func(t *testing.T) { - ctx := context.Background() - - shutdownFunc1 := func(ctx context.Context) error { - return errors.New("first shutdown error") - } - shutdownFunc2 := func(ctx context.Context) error { - return errors.New("second shutdown error") - } - - shutdownFuncs := []func(context.Context) error{ - shutdownFunc1, - shutdownFunc2, - } - - shutdown := shutdownOpenTelemetryComponents(shutdownFuncs) - - err := shutdown(ctx) - assert.Error(t, err) - assert.Equal(t, "second shutdown error", err.Error()) - }) -} - -func TestSRecordLatency(t *testing.T) { - cyx := context.Background() - srv := setupTestEndpoint(":7061", "/TestSRecordLatency") - time.Sleep(3 * time.Second) - defer func() { - assert.NoError(t, srv.Shutdown(cyx), "failed to shutdown srv") - }() - - var ds1 []func(context.Context) error - - cfg := &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7061/TestSRecordLatency", - } - - ot, ds, err := NewOpenTelemetry(cyx, cfg, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - - ot.RecordLatencyMetric(cyx, time.Now(), Attributes{Method: "handlePrepare"}) - assert.NoErrorf(t, err, "error occurred") - //when otel is disabled - cfg2 := &OTelConfig{ - OTELEnabled: false, - } - - ot1, ds, err2 := NewOpenTelemetry(cyx, cfg2, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - - ot1.RecordLatencyMetric(cyx, time.Now(), Attributes{Method: "handlePrepare"}) - - shutdownOpenTelemetryComponents(ds1) - assert.NoErrorf(t, err2, "error occurred") -} - -func TestRecordRequestCountMetric(t *testing.T) { - cyx := context.Background() - srv := setupTestEndpoint(":7061", "/TestRecordRequestCountMetric") - time.Sleep(3 * time.Second) - defer func() { - assert.NoError(t, srv.Shutdown(cyx), "failed to shutdown srv") - }() - - var ds1 []func(context.Context) error - - cfg := &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7061/TestRecordRequestCountMetric", - } - - ot, ds, err := NewOpenTelemetry(cyx, cfg, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - - ot.RecordRequestCountMetric(cyx, Attributes{Method: "handlePrepare"}) - - assert.NoErrorf(t, err, "error occurred") - - //when otel is disabled - cfg2 := &OTelConfig{ - OTELEnabled: false, - } - - ot1, ds, err2 := NewOpenTelemetry(cyx, cfg2, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - - ot1.RecordRequestCountMetric(cyx, Attributes{Method: "handlePrepare"}) - - shutdownOpenTelemetryComponents(ds1) - assert.NoErrorf(t, err2, "error occurred") -} - -func TestApplyTrace(t *testing.T) { - cyx := context.Background() - srv := setupTestEndpoint(":7061", "/TestApplyTrace") - time.Sleep(3 * time.Second) - defer func() { - assert.NoError(t, srv.Shutdown(cyx), "failed to shutdown srv") - }() - - var ds1 []func(context.Context) error - - cfg := &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7061/TestApplyTrace", - } - - ot, ds, err := NewOpenTelemetry(cyx, cfg, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - - _, span := ot.StartSpan(cyx, "test", []attribute.KeyValue{ - attribute.String("method", "handlePrepare"), - }) - - shutdownOpenTelemetryComponents(ds1) - assert.NoErrorf(t, err, "error occurred") - ot.EndSpan(span) -} - -func TestShutdown(t *testing.T) { - cyx := context.Background() - srv := setupTestEndpoint(":7061", "/TestShutdown") - time.Sleep(3 * time.Second) - defer func() { - assert.NoError(t, srv.Shutdown(cyx), "failed to shutdown srv") - }() - - var ds1 []func(context.Context) error - - cfg := &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7061/TestShutdown", - } - - _, ds, err := NewOpenTelemetry(cyx, cfg, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - shutdownOpenTelemetryComponents(ds1) -} - -// testHandler handles requests to the test endpoint. -func testHandler(w http.ResponseWriter, r *http.Request) { - fmt.Fprintln(w, "Hello, this is a test endpoint!") -} - -func setupTestEndpoint(st string, typ string) *http.Server { - // Create a new instance of a server - server := &http.Server{Addr: st} - - // Register the test handler with the specific type - http.HandleFunc(typ, testHandler) - - // Start the server in a goroutine - go func() { - if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { - fmt.Printf("ListenAndServe error: %s\n", err) - } - }() - - return server -} - -func TestSRecordError(t *testing.T) { - cyx := context.Background() - srv := setupTestEndpoint(":7061", "/TestSRecordError") - time.Sleep(3 * time.Second) - defer func() { - assert.NoError(t, srv.Shutdown(cyx), "failed to shutdown srv") - }() - - var ds1 []func(context.Context) error - - cfg := &OTelConfig{ - TracerEndpoint: "http://localhost:7060", - MetricEndpoint: "http://localhost:7061", - ServiceName: "test", - OTELEnabled: true, - TraceSampleRatio: 20, - Database: "testDB", - Instance: "testInstance", - HealthCheckEp: "localhost:7061/TestSRecordError", - } - - ot, ds, err := NewOpenTelemetry(cyx, cfg, nil) - ds1 = append(ds1, ds) - assert.NoErrorf(t, err, "error occurred") - - _, span := ot.StartSpan(cyx, "test", []attribute.KeyValue{ - attribute.String("method", "handlePrepare"), - }) - - ot.RecordError(span, fmt.Errorf("test error")) - - shutdownOpenTelemetryComponents(ds1) - assert.NoErrorf(t, err, "error occurred") -} - -func TestInitTracerProvider(t *testing.T) { - ctx := context.Background() - - tests := []struct { - name string - config *OTelConfig - wantErr bool - }{ - { - name: "Valid Configuration", - config: &OTelConfig{ - TracerEndpoint: "localhost:4317", - ServiceName: "test-service", - TraceSampleRatio: 1.0, - }, - wantErr: false, - }, - { - name: "Missing Tracer Endpoint", - config: &OTelConfig{ - TracerEndpoint: "", - ServiceName: "test-service", - TraceSampleRatio: 1.0, - }, - wantErr: true, - }, - { - name: "Invalid Tracer Endpoint Format", - config: &OTelConfig{ - TracerEndpoint: "invalid-endpoint", - ServiceName: "test-service", - TraceSampleRatio: 1.0, - }, - wantErr: true, - }, - { - name: "Zero Trace Sample Ratio", - config: &OTelConfig{ - TracerEndpoint: "localhost:4317", - ServiceName: "test-service", - TraceSampleRatio: 0.0, - }, - wantErr: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - resource := resource.NewSchemaless() // Mock resource - got, err := InitTracerProvider(ctx, tt.config, resource) - if (err != nil) != tt.wantErr { - t.Errorf("InitTracerProvider() error = %v, wantErr %v", err, tt.wantErr) - } - - if !tt.wantErr && (got == nil || reflect.TypeOf(got) != reflect.TypeOf(&sdktrace.TracerProvider{})) { - t.Errorf("InitTracerProvider() = %v, expected valid TracerProvider", got) - } - }) - } -} -func TestInitMeterProvider(t *testing.T) { - ctx := context.Background() - - tests := []struct { - name string - config *OTelConfig - wantErr bool - }{ - { - name: "Valid Configuration", - config: &OTelConfig{ - MetricEndpoint: "localhost:4318", - ServiceName: "test-service", - }, - wantErr: false, - }, - { - name: "Missing Metric Endpoint", - config: &OTelConfig{ - MetricEndpoint: "", - ServiceName: "test-service", - }, - wantErr: true, - }, - { - name: "Invalid Metric Endpoint Format", - config: &OTelConfig{ - MetricEndpoint: "://invalid-endpoint", - ServiceName: "test-service", - }, - wantErr: true, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - resource := resource.NewSchemaless() // Mock resource - got, err := InitMeterProvider(ctx, tt.config, resource) - - if (err != nil) != tt.wantErr { - t.Errorf("InitMeterProvider() error = %v, wantErr %v", err, tt.wantErr) - } - - if !tt.wantErr && (got == nil || reflect.TypeOf(got) != reflect.TypeOf(&sdkmetric.MeterProvider{})) { - t.Errorf("InitMeterProvider() = %v, expected valid MeterProvider", got) - } - }) - } -} - -func Test_buildOtelResource(t *testing.T) { - ctx := context.Background() - - tests := []struct { - name string - config *OTelConfig - }{ - { - name: "Valid Service Configuration", - config: &OTelConfig{ - ServiceName: "test-service", - ServiceVersion: "1.0.0", - }, - }, - { - name: "Missing Service Columns", - config: &OTelConfig{ - ServiceName: "", - ServiceVersion: "1.0.0", - }, - }, - { - name: "Missing Service Version", - config: &OTelConfig{ - ServiceName: "test-service", - ServiceVersion: "", - }, - }, - { - name: "Empty Config", - config: &OTelConfig{ - ServiceName: "", - ServiceVersion: "", - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := buildOtelResource(ctx, tt.config) - - // Validate that the returned resource is not nil - assert.NotNil(t, got, "buildOtelResource() returned nil for test case: %s", tt.name) - - // Retrieve attributes from resource - attrs := got.Attributes() - - // Check if expected attributes exist in the resource - expectedAttrs := map[string]string{ - string(semconv.ServiceNameKey): tt.config.ServiceName, - string(semconv.ServiceVersionKey): tt.config.ServiceVersion, - } - - for key, expectedValue := range expectedAttrs { - found := false - for _, attr := range attrs { - if string(attr.Key) == key && attr.Value.AsString() == expectedValue { - found = true - break - } - } - assert.True(t, found, "buildOtelResource() missing expected attribute: %s", key) - } - - // Ensure the instance ID exists - instanceIDFound := false - for _, attr := range attrs { - if string(attr.Key) == string(semconv.ServiceInstanceIDKey) { - instanceIDFound = true - break - } - } - assert.True(t, instanceIDFound, "buildOtelResource() missing expected attribute: ServiceInstanceID") - }) - } -} - -func TestOpenTelemetry_StartSpan(t *testing.T) { - ctx := context.Background() - - tests := []struct { - name string - otelConfig *OTelConfig - spanName string - attrs []attribute.KeyValue - expectSpan bool - }{ - { - name: "Start span when OTEL is enabled", - otelConfig: &OTelConfig{ - OTELEnabled: true, - }, - spanName: "test-span", - attrs: []attribute.KeyValue{attribute.String("key", "value")}, - expectSpan: true, - }, - { - name: "Return same context when OTEL is disabled", - otelConfig: &OTelConfig{ - OTELEnabled: false, - }, - spanName: "test-span", - attrs: []attribute.KeyValue{attribute.String("key", "value")}, - expectSpan: false, - }, - { - name: "Start span with empty attributes", - otelConfig: &OTelConfig{ - OTELEnabled: true, - }, - spanName: "empty-attributes-span", - attrs: []attribute.KeyValue{}, - expectSpan: true, - }, - { - name: "Start span with empty span name", - otelConfig: &OTelConfig{ - OTELEnabled: true, - }, - spanName: "", - attrs: []attribute.KeyValue{attribute.String("key", "value")}, - expectSpan: true, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - otelInstance := &OpenTelemetry{ - Config: tt.otelConfig, - tracer: noop.NewTracerProvider().Tracer("test"), // Use noop tracer for testing - logger: zap.NewNop(), - } - - newCtx, span := otelInstance.StartSpan(ctx, tt.spanName, tt.attrs) - - if tt.expectSpan { - assert.NotNil(t, span, "Expected a valid span but got nil") - } else { - assert.Nil(t, span, "Expected nil span but got a valid span") - } - - // Ensure context is updated when OTEL is enabled - if tt.expectSpan { - assert.NotEqual(t, ctx, newCtx, "Context should be updated with span") - } else { - assert.Equal(t, ctx, newCtx, "Context should remain unchanged when OTEL is disabled") - } - }) - } -} - -func TestOpenTelemetry_EndSpan(t *testing.T) { - tests := []struct { - name string - otelConfig *OTelConfig - expectCall bool - }{ - { - name: "End span when OTEL is enabled", - otelConfig: &OTelConfig{ - OTELEnabled: true, - }, - expectCall: true, - }, - { - name: "Do nothing when OTEL is disabled", - otelConfig: &OTelConfig{ - OTELEnabled: false, - }, - expectCall: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - otelInstance := &OpenTelemetry{ - Config: tt.otelConfig, - tracer: noop.NewTracerProvider().Tracer("test"), // No-op tracer - logger: zap.NewNop(), - } - - // Creating a valid context and span - ctx := context.Background() - otelInstance.tracer.Start(ctx, "test-span") - - // Spy to check if span.End() is called - ended := false - mockSpan := &mockSpan{endCalled: &ended} - - otelInstance.EndSpan(mockSpan) - - if tt.expectCall { - if !ended { - t.Errorf("Expected span.End() to be called, but it wasn't") - } - } else { - if ended { - t.Errorf("Expected span.End() to NOT be called, but it was") - } - } - }) - } -} - -// mockSpan is a mock implementation of trace.Span to check if End() is called. -type mockSpan struct { - trace.Span - endCalled *bool -} - -func (m *mockSpan) End(opts ...trace.SpanEndOption) { - *m.endCalled = true -} - -// MockOpenTelemetry is a mock implementation of OpenTelemetry for testing. -type MockOpenTelemetry struct { - *OpenTelemetry - RequestCountCalled bool - LatencyCalled bool -} - -// Override RecordMetrics to ensure the mock methods get called -func (m *MockOpenTelemetry) RecordMetrics(ctx context.Context, method string, startTime time.Time, queryType string, err error) { - // Check if OTEL is enabled before recording metrics - if !m.Config.OTELEnabled { - return - } - - // Call mock versions of RecordRequestCountMetric & RecordLatencyMetric - m.RecordRequestCountMetric(ctx, Attributes{Method: method, Status: "OK", QueryType: queryType}) - m.RecordLatencyMetric(ctx, startTime, Attributes{Method: method, QueryType: queryType}) -} - -// Override RecordRequestCountMetric to track calls -func (m *MockOpenTelemetry) RecordRequestCountMetric(ctx context.Context, attrs Attributes) { - m.RequestCountCalled = true -} - -// Override RecordLatencyMetric to track calls -func (m *MockOpenTelemetry) RecordLatencyMetric(ctx context.Context, startTime time.Time, attrs Attributes) { - m.LatencyCalled = true -} - -func TestOpenTelemetry_RecordMetrics(t *testing.T) { - tests := []struct { - name string - otelConfig *OTelConfig - method string - queryType string - err error - expectCall bool - }{ - { - name: "Record metrics with successful request", - otelConfig: &OTelConfig{ - OTELEnabled: true, - }, - method: "GET", - queryType: "select", - err: nil, - expectCall: true, - }, - { - name: "Record metrics with failed request", - otelConfig: &OTelConfig{ - OTELEnabled: true, - }, - method: "POST", - queryType: "insert", - err: errors.New("database error"), - expectCall: true, - }, - { - name: "Do not record metrics when OTEL is disabled", - otelConfig: &OTelConfig{ - OTELEnabled: false, - }, - method: "DELETE", - queryType: "delete", - err: nil, - expectCall: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() - startTime := time.Now() - - // Use MockOpenTelemetry with the overridden RecordMetrics method - mockOtel := &MockOpenTelemetry{ - OpenTelemetry: &OpenTelemetry{ - Config: tt.otelConfig, - tracer: noop.NewTracerProvider().Tracer("test"), - logger: zap.NewNop(), - }, - } - - // Call the function under test - mockOtel.RecordMetrics(ctx, tt.method, startTime, tt.queryType, tt.err) - - if tt.expectCall { - if !mockOtel.RequestCountCalled { - t.Errorf("Expected RecordRequestCountMetric to be called, but it wasn't") - } - if !mockOtel.LatencyCalled { - t.Errorf("Expected RecordLatencyMetric to be called, but it wasn't") - } - } else { - if mockOtel.RequestCountCalled || mockOtel.LatencyCalled { - t.Errorf("Expected no metric calls when OTEL is disabled, but they were called") - } - } - }) - } -} - -// MockSpan is a mock implementation of trace.Span for testing -type MockSpan struct { - trace.Span - events []string -} - -func (m *MockSpan) AddEvent(name string, opts ...trace.EventOption) { - m.events = append(m.events, name) -} - -func TestAddAnnotation(t *testing.T) { - tests := []struct { - name string - event string - }{ - { - name: "Add simple event", - event: "UserLoggedIn", - }, - { - name: "Add another event", - event: "DatabaseQueried", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - mockSpan := &MockSpan{} - ctx := trace.ContextWithSpan(context.Background(), mockSpan) - - // Call function - AddAnnotation(ctx, tt.event) - - // Verify that the event was recorded - assert.Contains(t, mockSpan.events, tt.event, "Expected event to be added") - }) - } -} - -func Test_isValidEndpoint(t *testing.T) { - type args struct { - endpoint string - } - tests := []struct { - name string - args args - want bool - }{ - { - name: "Valid host:port format", - args: args{endpoint: "localhost:8080"}, - want: true, - }, - { - name: "Invalid host:port format (no port)", - args: args{endpoint: "localhost"}, - want: false, - }, - { - name: "Valid URL format with http", - args: args{endpoint: "http://localhost:8080"}, - want: true, - }, - { - name: "Valid URL format with https", - args: args{endpoint: "https://example.com:443"}, - want: true, - }, - { - name: "Invalid URL format (missing host)", - args: args{endpoint: "http://:8080"}, - want: false, - }, - { - name: "Invalid URL format (missing port)", - args: args{endpoint: "http://localhost:"}, - want: false, - }, - { - name: "Empty string", - args: args{endpoint: ""}, - want: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := isValidEndpoint(tt.args.endpoint); got != tt.want { - t.Errorf("isValidEndpoint() = %v, want %v", got, tt.want) - } - }) - } -} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/tracing.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/tracing.go new file mode 100644 index 00000000..ec0f9190 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel/tracing.go @@ -0,0 +1,83 @@ +/* + * Copyright (C) 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy of + * the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ + +package otelgo + +import ( + "context" + "errors" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "go.opentelemetry.io/otel/trace/noop" + + texporter "github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/trace" +) + +type noOptionsSampler struct { + sampler sdktrace.Sampler +} + +func (s noOptionsSampler) ShouldSample(p sdktrace.SamplingParameters) sdktrace.SamplingResult { + if p.Name == "options" { + return sdktrace.SamplingResult{ + Decision: sdktrace.Drop, + Attributes: nil, + Tracestate: trace.SpanContextFromContext(p.ParentContext).TraceState(), + } + } + return s.sampler.ShouldSample(p) +} + +func (s noOptionsSampler) Description() string { + return "NoOptionsSampler(" + s.sampler.Description() + ")" +} + +func createTraceProvider(ctx context.Context, config *types.OtelConfig, res *resource.Resource) (trace.TracerProvider, ShutdownFn, error) { + if !config.Enabled { + return noop.NewTracerProvider(), func(_ context.Context) error { return nil }, nil + } + + sampler := sdktrace.Sampler(noOptionsSampler{sampler: sdktrace.ParentBased(sdktrace.TraceIDRatioBased(config.Traces.SamplingRatio))}) + + opts := []sdktrace.TracerProviderOption{ + sdktrace.WithResource(res), + sdktrace.WithSampler(sampler), + } + + if config.Traces.ProjectId != "" { + exporter, err := texporter.New(texporter.WithProjectID(config.Traces.ProjectId)) + if err != nil { + return nil, nil, err + } + opts = append(opts, sdktrace.WithBatcher(exporter)) + } else if config.Traces.Endpoint != "" { + if !isValidEndpoint(config.Traces.Endpoint) { + return nil, nil, errors.New("invalid Tracer endpoint format") + } + exporter, err := otlptracegrpc.New(ctx, otlptracegrpc.WithEndpoint(config.Traces.Endpoint), otlptracegrpc.WithInsecure()) + if err != nil { + return nil, nil, err + } + opts = append(opts, sdktrace.WithBatcher(exporter)) + } else { + return nil, nil, errors.New("no Tracer endpoint or project id provided") + } + result := sdktrace.NewTracerProvider(opts...) + return result, result.Shutdown, nil +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser/parser.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser/parser.go index 62815202..5aaf8bad 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser/parser.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser/parser.go @@ -2,9 +2,11 @@ package parser import ( "fmt" + "strings" + "sync" + cql "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/cqlparser" "github.com/antlr4-go/antlr/v4" - "strings" ) type syntaxErrorListener struct { @@ -29,129 +31,242 @@ func (l *syntaxErrorListener) SyntaxError( ) { // Format a clear error message with location l.errs = append(l.errs, msg) +} - recognizer.SetError(e) +type ProxyCqlParser struct { + p *cql.CqlParser + lexer *cql.CqlLexer + stream *antlr.CommonTokenStream + errorListener *syntaxErrorListener } -func NewParser(query string) *ProxyCqlParser { - errorListener := &syntaxErrorListener{ - DefaultErrorListener: antlr.NewDefaultErrorListener(), - errs: make([]string, 0), - } +// parserPool holds reusable parser instances to drastically reduce GC pressure and CPU overhead. +var parserPool = sync.Pool{ + New: func() interface{} { + // Pre-allocate slice capacity to prevent allocations during error appending + errorListener := &syntaxErrorListener{ + DefaultErrorListener: antlr.NewDefaultErrorListener(), + errs: make([]string, 0, 4), + } - lexer := cql.NewCqlLexer(antlr.NewInputStream(query)) - lexer.AddErrorListener(errorListener) + // Initialize with empty input to setup the structure once + input := antlr.NewInputStream("") + lexer := cql.NewCqlLexer(input) + lexer.RemoveErrorListeners() + lexer.AddErrorListener(errorListener) - stream := antlr.NewCommonTokenStream(lexer, antlr.TokenDefaultChannel) - parser := cql.NewCqlParser(stream) - parser.AddErrorListener(errorListener) + stream := antlr.NewCommonTokenStream(lexer, antlr.TokenDefaultChannel) + parser := cql.NewCqlParser(stream) + parser.RemoveErrorListeners() + parser.AddErrorListener(errorListener) + // By default, ANTLR tries to magically recover from bad queries but this is slow. Disable it. + parser.SetErrorHandler(antlr.NewBailErrorStrategy()) + // SLL ignores full context and runs exponentially faster, but occasionally fails on complex ambiguous syntax. + parser.GetInterpreter().SetPredictionMode(antlr.PredictionModeSLL) + return &ProxyCqlParser{ + p: parser, + lexer: lexer, + stream: stream, + errorListener: errorListener, + } + }, +} - return &ProxyCqlParser{ - p: parser, - errorListener: errorListener, - } +// GetParser pulls a parser from the pool and resets its state for the new query. +func GetParser(query string) *ProxyCqlParser { + proxy := parserPool.Get().(*ProxyCqlParser) + + // Clear the errors without reallocating the underlying array + proxy.errorListener.errs = proxy.errorListener.errs[:0] + + // Create a new input stream for the query and update the existing objects + input := antlr.NewInputStream(query) + proxy.lexer.SetInputStream(input) + proxy.stream.SetTokenSource(proxy.lexer) + proxy.p.SetTokenStream(proxy.stream) + + return proxy } -type ProxyCqlParser struct { - p *cql.CqlParser - errorListener *syntaxErrorListener +// Release returns the parser back to the pool. MUST be called after parsing. +func (p *ProxyCqlParser) Release() { + parserPool.Put(p) } +// --- Wrapper Methods --- + func (p *ProxyCqlParser) ValidateNoErrors() error { - if len(p.errorListener.errs) > 0 { - return fmt.Errorf("parsing error: %s", strings.Join(p.errorListener.errs, ", ")) - } - return nil + return p.errorListener.ValidateNoErrors() } -func (p *ProxyCqlParser) AlterTable() (cql.IAlterTableContext, error) { - result := p.p.AlterTable() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) AlterTable() (res cql.IAlterTableContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.AlterTable() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) Delete_() (cql.IDelete_Context, error) { - result := p.p.Delete_() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) Delete_() (res cql.IDelete_Context, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.Delete_() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) Use_() (cql.IUse_Context, error) { - result := p.p.Use_() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) Use_() (res cql.IUse_Context, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.Use_() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) Select_() (cql.ISelect_Context, error) { - result := p.p.Select_() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) Select_() (res cql.ISelect_Context, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.Select_() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) Truncate() (cql.ITruncateContext, error) { - result := p.p.Truncate() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) Truncate() (res cql.ITruncateContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.Truncate() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) DescribeStatement() (cql.IDescribeStatementContext, error) { - result := p.p.DescribeStatement() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) DescribeStatement() (res cql.IDescribeStatementContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.DescribeStatement() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) CreateTable() (cql.ICreateTableContext, error) { - result := p.p.CreateTable() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) CreateTable() (res cql.ICreateTableContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.CreateTable() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) DropTable() (cql.IDropTableContext, error) { - result := p.p.DropTable() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) DropTable() (res cql.IDropTableContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.DropTable() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) Update() (cql.IUpdateContext, error) { - result := p.p.Update() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) Update() (res cql.IUpdateContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.Update() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } -func (p *ProxyCqlParser) Insert() (cql.IInsertContext, error) { - result := p.p.Insert() - err := p.ValidateNoErrors() - if err != nil { +func (p *ProxyCqlParser) Insert() (res cql.IInsertContext, err error) { + defer func() { + if r := recover(); r != nil { + if e := p.ValidateNoErrors(); e != nil { + err = e + } else { + err = fmt.Errorf("parsing error: %T", r) + } + } + }() + res = p.p.Insert() + if err := p.ValidateNoErrors(); err != nil { return nil, err } - return result, nil + return res, nil } func (p *ProxyCqlParser) GetFirstToken() antlr.Token { diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/proxy.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/proxy.go index 0af645ab..346166a6 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/proxy.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/proxy.go @@ -17,6 +17,8 @@ package main import ( "context" "fmt" + "net/http" + _ "net/http/pprof" // This side effect import registers the handlers "os" "os/signal" @@ -27,6 +29,10 @@ func main() { ctx, cancel := signalContext(context.Background(), os.Interrupt, os.Kill) defer cancel() + go func() { + http.ListenAndServe("localhost:6060", nil) + }() + err := proxy.Run(ctx, os.Args[1:]) if err != nil { fmt.Fprintf(os.Stderr, "error: %v\n", err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_batch.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_batch.go new file mode 100644 index 00000000..864cbf2a --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_batch.go @@ -0,0 +1,108 @@ +// Copyright (c) DataStax, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package request_handlers + +import ( + "context" + "errors" + "fmt" + bigtableModule "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.opentelemetry.io/otel/trace" + "strings" +) + +type BatchRequestHandler struct { + server IProxyServer +} + +func (b *BatchRequestHandler) Name() string { + return "batch" +} + +func (b *BatchRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodeBatch +} + +func (b *BatchRequestHandler) HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + msg := req.msg.(*proxy_types.PartialBatch) + span := trace.SpanFromContext(ctx) + + bulkMutations, keyspace, err := b.bindBulkOperations(msg, session, req.header.Version) + if err != nil { + return &message.ConfigError{ErrorMessage: err.Error()}, err + } + span.AddEvent(proxy_types.SendingBulkApplyMutation) + var errs []string + for tableName, mutations := range bulkMutations.Mutations() { + res, err := b.server.BigtableClient().ApplyBulkMutation(ctx, keyspace, tableName, mutations) + if err != nil { + errs = append(errs, err.Error()) + } else if res.FailedRows != "" { + err = fmt.Errorf("failed rows for table %s: %s", tableName, res.FailedRows) + errs = append(errs, res.FailedRows) + } + } + span.AddEvent(proxy_types.GotBulkApplyResp) + span.SetAttributes() + if len(errs) > 0 { + err := errors.New(strings.Join(errs, "\n")) + return &message.ServerError{ErrorMessage: err.Error()}, err + } + + return &message.VoidResult{}, nil +} + +func (b *BatchRequestHandler) bindBulkOperations(msg *proxy_types.PartialBatch, session IProxySession, pv primitive.ProtocolVersion) (*bigtableModule.BigtableBulkMutation, types.Keyspace, error) { + var keyspace types.Keyspace + tableMutationsMap := bigtableModule.NewBigtableBulkMutation() + for index, queryId := range msg.QueryOrIds { + queryOrId, ok := queryId.([]byte) + if !ok { + return nil, "", fmt.Errorf("batch query id malformed") + } + id := proxy_types.PreparedIdKey(queryOrId) + preparedStmt, ok := b.server.PreparedQueryCache().Load(id) + if !ok { + return nil, "", fmt.Errorf("prepared query not found in cache") + } + + if preparedStmt.Keyspace() != "" { + keyspace = preparedStmt.Keyspace() + } + + // note: we don't support batch named queries at this time + executableQuery, err := b.server.Translator().BindQuery(preparedStmt, msg.BatchPositionalValues[index], nil, pv) + if err != nil { + return nil, "", err + } + mutation, ok := executableQuery.AsBulkMutation() + if !ok { + return nil, "", fmt.Errorf("query type '%s' not compatible with bulk", executableQuery.QueryType().String()) + } + tableMutationsMap.AddMutation(mutation) + } + if keyspace == "" { + keyspace = session.SessionKeyspace() + } + return tableMutationsMap, keyspace, nil +} + +func NewBatchRequestHandler(server IProxyServer) IProxyRequestHandler { + return &BatchRequestHandler{server: server} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_execute.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_execute.go new file mode 100644 index 00000000..29c8fbc9 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_execute.go @@ -0,0 +1,77 @@ +// Copyright (c) DataStax, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package request_handlers + +import ( + "context" + "errors" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.opentelemetry.io/otel/trace" +) + +type ExecuteRequestHandler struct { + server IProxyServer +} + +func (e *ExecuteRequestHandler) Name() string { + return "execute" +} + +func (e *ExecuteRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodeExecute +} + +func (e *ExecuteRequestHandler) HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + msg := req.msg.(*proxy_types.PartialExecute) + span := trace.SpanFromContext(ctx) + + var preparedStmt types.IPreparedQuery + + id := proxy_types.PreparedIdKey(msg.QueryId) + + var ok bool + preparedStmt, ok = e.server.PreparedQueryCache().Load(id) + if !ok { + err := errors.New(proxy_types.ErrQueryNotPrepared) + return &message.ServerError{ErrorMessage: proxy_types.ErrQueryNotPrepared}, err + } + + otelgo.AddQueryAnnotations(span, preparedStmt) + + span.AddEvent("bind-query") + boundQuery, err := e.server.Translator().BindQuery(preparedStmt, msg.PositionalValues, msg.NamedValues, req.header.Version) + if err != nil { + return &message.ConfigError{ErrorMessage: err.Error()}, err + } + + span.AddEvent("execute-query") + results, err := e.server.Executor().Execute(ctx, session, boundQuery) + if err != nil { + return &message.ConfigError{ErrorMessage: err.Error()}, err + } + if preparedStmt.QueryType().IsDDLType() { + span.AddEvent("execute-ddl-event") + e.server.HandlePostDDLEvent(preparedStmt.QueryType(), preparedStmt.Keyspace(), preparedStmt.Table()) + } + return results, nil +} + +func NewExecuteRequestHandler(server IProxyServer) IProxyRequestHandler { + return &ExecuteRequestHandler{server: server} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_options.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_options.go new file mode 100644 index 00000000..10868384 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_options.go @@ -0,0 +1,30 @@ +package request_handlers + +import ( + "context" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" +) + +type OptionsRequestHandler struct { + server IProxyServer +} + +func (o *OptionsRequestHandler) Name() string { + return "options" +} + +func (o *OptionsRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodeOptions +} + +func (o *OptionsRequestHandler) HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + return &message.Supported{Options: map[string][]string{ + "CQL_VERSION": {o.server.CQLVersion()}, + "COMPRESSION": {}, + }}, nil +} + +func NewOptionsRequestHandler(server IProxyServer) IProxyRequestHandler { + return &OptionsRequestHandler{server: server} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_prepare.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_prepare.go new file mode 100644 index 00000000..8a65571e --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_prepare.go @@ -0,0 +1,116 @@ +// Copyright (c) DataStax, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package request_handlers + +import ( + "context" + "fmt" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler" + cql "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/cqlparser" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.opentelemetry.io/otel/trace" + "go.uber.org/zap" +) + +type PrepareRequestHandler struct { + server IProxyServer +} + +func (p *PrepareRequestHandler) Name() string { + return "prepare" +} + +func (p *PrepareRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodePrepare +} +func (p *PrepareRequestHandler) HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + msg := req.msg.(*message.Prepare) + span := trace.SpanFromContext(ctx) + + qt := types.QueryTypeUnknown + + id := getQueryId(session, msg) + if preparedQuery, found := p.server.PreparedQueryCache().Load(id); found { + qt = preparedQuery.QueryType() + return responsehandler.BuildPreparedResultResponse(id, preparedQuery), nil + } + + p.server.Logger().Debug("preparing query", zap.String(proxy_types.Query, msg.Query), zap.Int16("stream", req.header.StreamId)) + + keyspace := session.SessionKeyspace() + if len(msg.Keyspace) != 0 { + keyspace = types.Keyspace(msg.Keyspace) + } + + pParser := parser.GetParser(msg.Query) + var err error + qt, err = parseQueryType(pParser) + if err != nil { + p.server.Logger().Error("failed to parse query type", zap.Error(err), zap.String("cql", msg.Query)) + return &message.Invalid{ErrorMessage: err.Error()}, err + } + + rawQuery := types.NewRawQuery(req.header, keyspace, msg.Query, pParser, qt) + + preparedQuery, err := prepareQuery(ctx, p.server, session, rawQuery) + if err != nil { + return nil, err + } + otelgo.AddQueryAnnotations(span, preparedQuery) + + // update query cache + p.server.PreparedQueryCache().Store(id, preparedQuery) + + resp := responsehandler.BuildPreparedResultResponse(id, preparedQuery) + return resp, nil +} + +func NewPrepareRequestHandler(server IProxyServer) IProxyRequestHandler { + return &PrepareRequestHandler{server: server} +} + +func parseQueryType(p *parser.ProxyCqlParser) (types.QueryType, error) { + tok := p.GetFirstToken() + t := tok.GetTokenType() + switch t { + case cql.CqlLexerK_SELECT: + return types.QueryTypeSelect, nil + case cql.CqlLexerK_INSERT: + return types.QueryTypeInsert, nil + case cql.CqlLexerK_UPDATE: + return types.QueryTypeUpdate, nil + case cql.CqlLexerK_DELETE: + return types.QueryTypeDelete, nil + case cql.CqlLexerK_CREATE: + return types.QueryTypeCreate, nil + case cql.CqlLexerK_ALTER: + return types.QueryTypeAlter, nil + case cql.CqlLexerK_DROP: + return types.QueryTypeDrop, nil + case cql.CqlLexerK_TRUNCATE: + return types.QueryTypeTruncate, nil + case cql.CqlLexerK_USE: + return types.QueryTypeUse, nil + case cql.CqlLexerK_DESCRIBE, cql.CqlLexerK_DESC: + return types.QueryTypeDescribe, nil + default: + return types.QueryTypeUnknown, fmt.Errorf("unsupported query type: %s", tok.String()) + } +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_query.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_query.go new file mode 100644 index 00000000..9ffd4dff --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_query.go @@ -0,0 +1,86 @@ +// Copyright (c) DataStax, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package request_handlers + +import ( + "context" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" + "go.uber.org/zap" + "time" +) + +type QueryRequestHandler struct { + server IProxyServer +} + +func (q *QueryRequestHandler) Name() string { + return "query" +} + +func (q *QueryRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodeQuery +} + +func (q *QueryRequestHandler) HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + msg := req.msg.(*proxy_types.PartialQuery) + span := trace.SpanFromContext(ctx) + q.server.Logger().Debug("handling query", zap.String("encodedQuery", msg.Query), zap.Int16("stream", req.header.StreamId)) + + p := parser.GetParser(msg.Query) + qt, err := parseQueryType(p) + if err != nil { + return &message.Invalid{ErrorMessage: err.Error()}, err + } + if span.IsRecording() { + span.SetAttributes(attribute.String(proxy_types.QueryTypeConst, qt.String())) + } + + rawQuery := types.NewRawQuery(req.header, session.SessionKeyspace(), msg.Query, p, qt) + + query, err := prepareQuery(ctx, q.server, session, rawQuery) + if err != nil { + return &message.ServerError{ErrorMessage: err.Error()}, err + } + + values := types.NewQueryParameterValues(query.Parameters(), time.Now()) + executableQuery, err := q.server.Translator().BindQueryParameters(query, values, req.header.Version) + if err != nil { + return &message.ConfigError{ErrorMessage: err.Error()}, err + } + + span.AddEvent(proxy_types.ExecutingBigtableSQLAPIRequestEvent) + selectResult, err := q.server.Executor().Execute(ctx, session, executableQuery) + span.AddEvent(proxy_types.BigtableExecutionDoneEvent) + + if err != nil { + return &message.ServerError{ErrorMessage: err.Error()}, err + } + + if rawQuery.QueryType().IsDDLType() { + q.server.HandlePostDDLEvent(query.QueryType(), query.Keyspace(), query.Table()) + } + + return selectResult, nil +} + +func NewQueryRequestHandler(server IProxyServer) IProxyRequestHandler { + return &QueryRequestHandler{server: server} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_register.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_register.go new file mode 100644 index 00000000..c27452fb --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_register.go @@ -0,0 +1,32 @@ +package request_handlers + +import ( + "context" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" +) + +type RegisterRequestHandler struct { +} + +func (r *RegisterRequestHandler) Name() string { + return "register" +} + +func (r *RegisterRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodeRegister +} + +func (r *RegisterRequestHandler) HandleRequest(_ context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + msg := req.msg.(*message.Register) + for _, t := range msg.EventTypes { + if t == primitive.EventTypeSchemaChange { + session.RegisterForEvents() + } + } + return &message.Ready{}, nil +} + +func NewRegisterRequestHandler() IProxyRequestHandler { + return &RegisterRequestHandler{} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_startup.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_startup.go new file mode 100644 index 00000000..5bd52a15 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handle_startup.go @@ -0,0 +1,27 @@ +package request_handlers + +import ( + "context" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" +) + +type StartupRequestHandler struct { +} + +func (s *StartupRequestHandler) Name() string { + return "startup" +} + +func (s *StartupRequestHandler) OpCode() primitive.OpCode { + return primitive.OpCodeStartup +} + +func (s *StartupRequestHandler) HandleRequest(_ context.Context, _ IProxySession, req *ProxyRequest) (message.Message, error) { + // CC - register for Event types and respond READY + return &message.Ready{}, nil +} + +func NewStartupRequestHandler() IProxyRequestHandler { + return &StartupRequestHandler{} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handler_manager.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handler_manager.go new file mode 100644 index 00000000..5a0bfb95 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/handler_manager.go @@ -0,0 +1,84 @@ +package request_handlers + +import ( + "context" + "fmt" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" +) + +type HandlerManager struct { + requestHandlers map[primitive.OpCode]IProxyRequestHandler +} + +func NewHandlerManager() *HandlerManager { + return &HandlerManager{ + requestHandlers: make(map[primitive.OpCode]IProxyRequestHandler), + } +} + +func (h *HandlerManager) InitHandlers(server IProxyServer) { + handlers := []IProxyRequestHandler{ + NewOptionsRequestHandler(server), + NewStartupRequestHandler(), + NewRegisterRequestHandler(), + NewPrepareRequestHandler(server), + NewExecuteRequestHandler(server), + NewQueryRequestHandler(server), + NewBatchRequestHandler(server), + } + h.requestHandlers = make(map[primitive.OpCode]IProxyRequestHandler) + for _, handler := range handlers { + h.requestHandlers[handler.OpCode()] = handler + } +} + +func (h *HandlerManager) HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) { + req.Attributes.Method = opCodeToNiceString(req.header.OpCode) + handler, ok := h.requestHandlers[req.header.OpCode] + if !ok { + return nil, fmt.Errorf("unsupported operation: %s", opCodeToNiceString(req.header.OpCode)) + } + return handler.HandleRequest(ctx, session, req) +} + +func opCodeToNiceString(c primitive.OpCode) string { + switch c { + case primitive.OpCodeStartup: + return "STARTUP" + case primitive.OpCodeOptions: + return "OPTIONS" + case primitive.OpCodeQuery: + return "QUERY" + case primitive.OpCodePrepare: + return "PREPARE" + case primitive.OpCodeExecute: + return "EXECUTE" + case primitive.OpCodeRegister: + return "REGISTER" + case primitive.OpCodeBatch: + return "BATCH" + case primitive.OpCodeAuthResponse: + return "AUTH RESPONSE" + case primitive.OpCodeDseRevise: + return "REVISE" + // responses + case primitive.OpCodeError: + return "ERROR" + case primitive.OpCodeReady: + return "READY" + case primitive.OpCodeAuthenticate: + return "AUTHENTICATE" + case primitive.OpCodeSupported: + return "SUPPORTED" + case primitive.OpCodeResult: + return "RESULT" + case primitive.OpCodeEvent: + return "EVENT" + case primitive.OpCodeAuthChallenge: + return "AUTH CHALLENGE" + case primitive.OpCodeAuthSuccess: + return "AUTH SUCCESS" + } + return "UNKNOWN" +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/types.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/types.go new file mode 100644 index 00000000..7ecb7ef0 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/types.go @@ -0,0 +1,46 @@ +package request_handlers + +import ( + "context" + bigtableModule "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators" + "github.com/datastax/go-cassandra-native-protocol/frame" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.uber.org/zap" +) + +type IProxyRequestHandler interface { + Name() string + OpCode() primitive.OpCode + HandleRequest(ctx context.Context, session IProxySession, req *ProxyRequest) (message.Message, error) +} + +type IProxySession interface { + types.ICassandraClient + SessionKeyspace() types.Keyspace + RegisterForEvents() +} + +type IProxyServer interface { + CQLVersion() string + PreparedQueryCache() proxycore.PreparedCache[types.IPreparedQuery] + BigtableClient() *bigtableModule.BigtableAdapter + Translator() *translators.TranslatorManager + Executor() *executors.QueryExecutorManager + HandlePostDDLEvent(queryType types.QueryType, keyspace types.Keyspace, table types.TableName) + Logger() *zap.Logger +} + +type ProxyRequest struct { + header *frame.Header + msg message.Message + Attributes types.Attributes +} + +func NewProxyRequest(header *frame.Header, msg message.Message) *ProxyRequest { + return &ProxyRequest{header: header, msg: msg, Attributes: types.Attributes{}} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/utils.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/utils.go new file mode 100644 index 00000000..97f4fcc8 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers/utils.go @@ -0,0 +1,28 @@ +package request_handlers + +import ( + "context" + "crypto/md5" + "fmt" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + "github.com/datastax/go-cassandra-native-protocol/message" +) + +func getQueryId(session IProxySession, msg *message.Prepare) [16]byte { + return md5.Sum([]byte(msg.Query + string(session.SessionKeyspace()))) +} + +func prepareQuery(ctx context.Context, server IProxyServer, session IProxySession, query *types.RawQuery) (types.IPreparedQuery, error) { + preparedQuery, err := server.Translator().TranslateQuery(ctx, query, session.SessionKeyspace()) + if err != nil { + return nil, err + } + + btPreparedQuery, err := server.BigtableClient().PrepareStatement(ctx, preparedQuery) + if err != nil { + return nil, fmt.Errorf("failed to prepare bigtable statement `%s`: %w", preparedQuery.BigtableQuery(), err) + } + preparedQuery.SetBigtablePreparedQuery(btPreparedQuery) + + return preparedQuery, err +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler/responsehandler_utils.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler/responsehandler_utils.go index 3fe9a17b..195c94dd 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler/responsehandler_utils.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler/responsehandler_utils.go @@ -53,7 +53,7 @@ func BuildRowsResultResponse(st *types.ExecutableSelectQuery, rows []types.GoRow }, nil } -func BuildPreparedResultResponse(id [16]byte, query types.IPreparedQuery) (*message.PreparedResult, error) { +func BuildPreparedResultResponse(id [16]byte, query types.IPreparedQuery) *message.PreparedResult { resultColumns := query.ResponseColumns() var resultMetadata *message.RowsMetadata = nil if query.QueryType() == types.QueryTypeSelect { @@ -70,7 +70,7 @@ func BuildPreparedResultResponse(id [16]byte, query types.IPreparedQuery) (*mess PkIndices: pkIndices, Columns: variableMetadata, }, - }, nil + } } func buildPreparedResult(query types.IPreparedQuery) ([]*message.ColumnMetadata, []uint16) { diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/client.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/client.go new file mode 100644 index 00000000..30ab377b --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/client.go @@ -0,0 +1,156 @@ +// Copyright (c) DataStax, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package proxy + +import ( + "bytes" + "context" + "errors" + "fmt" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore" + "github.com/datastax/go-cassandra-native-protocol/frame" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.opentelemetry.io/otel/codes" + "go.uber.org/zap" + "io" + "runtime/debug" +) + +type Sender interface { + Send(hdr *frame.Header, msg message.Message) +} + +type client struct { + sessionKeyspace types.Keyspace + proxy *Server + conn *proxycore.Conn + sender Sender +} + +func (c *client) SetSessionKeyspace(k types.Keyspace) { + c.sessionKeyspace = k +} + +func (c *client) SessionKeyspace() types.Keyspace { + return c.sessionKeyspace +} + +func (c *client) Proxy() proxy_types.IProxy { + return c.proxy +} + +func (c *client) RegisterForEvents() { + c.proxy.registerForEvents(c) +} + +func (c *client) Receive(reader io.Reader) error { + ctx, cancel := context.WithCancel(c.proxy.ctx) + defer cancel() + + otelCtx, span := c.proxy.tracer.Start(ctx, "receive") + defer span.End() + span.AddEvent("decode-raw-request") + raw, err := codec.DecodeRawFrame(reader) + if err != nil { + if !errors.Is(err, io.EOF) { + c.proxy.logger.Error("unable to decode frame", zap.Error(err)) + } + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return err + } + + defer func() { + if r := recover(); r != nil { + fmt.Printf("Trapped a panic: %v\n", r) + + fmt.Println("Stack Trace:") + debug.PrintStack() + + c.sender.Send(raw.Header, &message.ProtocolError{ + ErrorMessage: "internal error", + }) + + span.SetStatus(codes.Error, "internal error") + } + }() + + if raw.Header.Version > c.proxy.config.Options.MaxProtocolVersion || raw.Header.Version < primitive.ProtocolVersion3 { + // IMPORTANT - do not change this message - it's parsed by CQL clients when negotiating the protocol version + errorMessage := fmt.Sprintf("Invalid or unsupported protocol version %d", raw.Header.Version) + c.sender.Send(raw.Header, &message.ProtocolError{ + ErrorMessage: errorMessage, + }) + err := errors.New(errorMessage) + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return nil + } + + span.AddEvent("decode-request-body") + body, err := codec.DecodeBody(raw.Header, bytes.NewReader(raw.Body)) + if err != nil { + c.proxy.logger.Error("unable to decode body", zap.Error(err)) + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return err + } + + span.AddEvent("handle-request") + req := request_handlers.NewProxyRequest(raw.Header, body.Message) + response, err := c.proxy.handlerManager.HandleRequest(otelCtx, c, req) + if err != nil { + span.SetAttributes(otelgo.CommonAttributes(req.Attributes)...) + c.proxy.logger.Error("error handling request", zap.Error(err)) + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return err + } + + span.SetAttributes(otelgo.CommonAttributes(req.Attributes)...) + span.AddEvent("send-response") + c.sender.Send(raw.Header, response) + span.AddEvent("done") + span.SetStatus(codes.Ok, "") + return nil +} + +func (c *client) Send(hdr *frame.Header, msg message.Message) { + _ = c.conn.Write(proxycore.SenderFunc(func(writer io.Writer) error { + return codec.EncodeFrame(frame.NewFrame(hdr.Version, hdr.StreamId, msg), writer) + })) +} + +func (c *client) Closing(_ error) { + c.proxy.removeClient(c) +} + +// handleEvent handles events from the proxy core +// It sends the event message to all connected clients. +func (c *client) handleEvent(event proxycore.Event) { + switch evt := event.(type) { + case *proxycore.SchemaChangeEvent: + c.sender.Send(&frame.Header{ + Version: c.proxy.config.Options.ProtocolVersion, + StreamId: -1, // -1 for events + OpCode: primitive.OpCodeEvent, + }, evt.Message) + } +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs.go index fccab759..d4af2fb6 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs.go @@ -15,9 +15,9 @@ package proxy import ( - "encoding/hex" "errors" "fmt" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" "io" "github.com/datastax/go-cassandra-native-protocol/frame" @@ -41,7 +41,7 @@ func (c *partialQueryCodec) Decode(source io.Reader, _ primitive.ProtocolVersion if query, err := primitive.ReadLongString(source); err != nil { return nil, err } else { - return &partialQuery{query}, nil + return &proxy_types.PartialQuery{Query: query}, nil } } @@ -49,56 +49,6 @@ func (c *partialQueryCodec) GetOpCode() primitive.OpCode { return primitive.OpCodeQuery } -type partialQuery struct { - query string -} - -func (p *partialQuery) IsResponse() bool { - return false -} - -func (p *partialQuery) GetOpCode() primitive.OpCode { - return primitive.OpCodeQuery -} - -func (p *partialQuery) Clone() message.Message { - return &partialQuery{p.query} -} - -type partialExecute struct { - queryId []byte - // The positional values of the associated query. Positional values are designated in query strings with a - // question mark bind marker ('?'). Positional values are valid for all protocol versions. - // It is illegal to use both positional and named values at the same time. If this happens, positional values will - // be used and named values will be silently ignored. - PositionalValues []*primitive.Value - - // The named values of the associated query. Named values are designated in query strings with a - // a bind marker starting with a colon (e.g. ':var1'). Named values can only be used with protocol version 3 or - // higher. - // It is illegal to use both positional and named values at the same time. If this happens, positional values will - // be used and named values will be silently ignored. - NamedValues map[string]*primitive.Value -} - -func (m *partialExecute) IsResponse() bool { - return false -} - -func (m *partialExecute) GetOpCode() primitive.OpCode { - return primitive.OpCodeExecute -} - -func (m *partialExecute) Clone() message.Message { - return &partialExecute{ - queryId: primitive.CloneByteSlice(m.queryId), - } -} - -func (m *partialExecute) String() string { - return "EXECUTE " + hex.EncodeToString(m.queryId) -} - type partialExecuteCodec struct{} func (c *partialExecuteCodec) Encode(_ message.Message, _ io.Writer, _ primitive.ProtocolVersion) error { @@ -110,10 +60,10 @@ func (c *partialExecuteCodec) EncodedLength(_ message.Message, _ primitive.Proto } func (c *partialExecuteCodec) Decode(source io.Reader, protocolV primitive.ProtocolVersion) (msg message.Message, err error) { - execute := &partialExecute{} - if execute.queryId, err = primitive.ReadShortBytes(source); err != nil { + execute := &proxy_types.PartialExecute{} + if execute.QueryId, err = primitive.ReadShortBytes(source); err != nil { return nil, fmt.Errorf("cannot read EXECUTE query id: %w", err) - } else if len(execute.queryId) == 0 { + } else if len(execute.QueryId) == 0 { return nil, errors.New("EXECUTE missing query id") } options, err := message.DecodeQueryOptions(source, protocolV) @@ -129,31 +79,6 @@ func (c *partialExecuteCodec) GetOpCode() primitive.OpCode { return primitive.OpCodeExecute } -type partialBatch struct { - queryOrIds []interface{} - // The positional values of the associated query. Positional values are designated in query strings with a - // question mark bind marker ('?'). Positional values are valid for all protocol versions. - // It is illegal to use both positional and named values at the same time. If this happens, positional values will - // be used and named values will be silently ignored. - BatchPositionalValues [][]*primitive.Value -} - -func (p partialBatch) IsResponse() bool { - return false -} - -func (p partialBatch) GetOpCode() primitive.OpCode { - return primitive.OpCodeBatch -} - -func (p partialBatch) Clone() message.Message { - queryOrIds := make([]interface{}, len(p.queryOrIds)) - copy(queryOrIds, p.queryOrIds) - PositionalValues := make([][]*primitive.Value, len(p.BatchPositionalValues)) - copy(PositionalValues, p.BatchPositionalValues) - return &partialBatch{queryOrIds, PositionalValues} -} - type partialBatchCodec struct{} func (p partialBatchCodec) Encode(msg message.Message, dest io.Writer, version primitive.ProtocolVersion) error { @@ -165,7 +90,7 @@ func (p partialBatchCodec) EncodedLength(msg message.Message, version primitive. } func (p partialBatchCodec) Decode(source io.Reader, version primitive.ProtocolVersion) (msg message.Message, err error) { - var partialBatchExecute = &partialBatch{} + var partialBatchExecute = &proxy_types.PartialBatch{} var positionalValues []*primitive.Value var queryOrIds []interface{} var typ uint8 @@ -205,7 +130,7 @@ func (p partialBatchCodec) Decode(source io.Reader, version primitive.ProtocolVe } queryOrIds[i] = queryOrId } - partialBatchExecute.queryOrIds = queryOrIds + partialBatchExecute.QueryOrIds = queryOrIds return partialBatchExecute, nil } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs_test.go index fcbea9a9..a2bec4d5 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/codecs_test.go @@ -17,6 +17,7 @@ package proxy import ( "bytes" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types" "io" "testing" @@ -25,7 +26,7 @@ import ( ) func Test_IsResponsePartialQuery(t *testing.T) { - var query partialQuery + var query proxy_types.PartialQuery res := query.IsResponse() assert.Equalf(t, res, false, "required to be false") @@ -37,7 +38,7 @@ func Test_IsResponsePartialQuery(t *testing.T) { } func Test_IsResponsePartialExecute(t *testing.T) { - var query partialExecute + var query proxy_types.PartialExecute query.IsResponse() } @@ -82,14 +83,14 @@ func Test_EncodepartialQueryCodecEncodedLength(t *testing.T) { assert.Errorf(t, err, "function should return error") } func TestPartialExecute(t *testing.T) { - pe := partialExecute{} + pe := proxy_types.PartialExecute{} res := pe.GetOpCode() assert.NotNilf(t, res, "function should return not nil") } func TestPartialExecuteClone(t *testing.T) { - pe := partialExecute{ - queryId: []byte("testqueryid"), + pe := proxy_types.PartialExecute{ + QueryId: []byte("testqueryid"), } res := pe.Clone() assert.NotNilf(t, res, "should not be nil") @@ -112,8 +113,7 @@ func TestPartialExecuteCodecDecode(t *testing.T) { } func TestPartialBatchClone(t *testing.T) { - - pa := partialBatch{} + pa := proxy_types.PartialBatch{} res := pa.Clone() assert.NotNilf(t, res, "should not be nil") } @@ -126,8 +126,8 @@ func TestPartialQueryCodecDecode(t *testing.T) { } func Test_partialExecuteString(t *testing.T) { - a := partialExecute{ - queryId: []byte("test"), + a := proxy_types.PartialExecute{ + QueryId: []byte("test"), } res := a.String() assert.NotNilf(t, res, "should not be nil") @@ -152,7 +152,7 @@ func Test_partialExecuteCodecEncode(t *testing.T) { a.Encode(nil, writer, primitive.ProtocolVersion3) } -func Test_partialExecuteCodecDecode(t *testing.T) { +func Test_partialExecuteCodecEncodedLength(t *testing.T) { a := partialExecuteCodec{} var writer io.Writer @@ -172,13 +172,13 @@ func Test_partialExecuteCodecDecode(t *testing.T) { } func Test_partialBatchIsResponse(t *testing.T) { - p := partialBatch{} + p := proxy_types.PartialBatch{} res := p.IsResponse() assert.Equalf(t, false, res, "should return falses") } func Test_partialBatchGetOpCode(t *testing.T) { - p := partialBatch{} + p := proxy_types.PartialBatch{} res := p.GetOpCode() assert.NotNilf(t, res, "should not return nil") } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/defaults.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/defaults.go index 847806f9..7a25bfe1 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/defaults.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/defaults.go @@ -24,7 +24,7 @@ import ( var ( // todo ensure this is a reasonable default - DefaultBigtableGrpcChannels = 1 + DefaultBigtableGrpcChannels = 4 BigtableMinSession = 100 BigtableMaxSession = 400 DefaultSchemaMappingTableName = "schema_mapping" diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup.go index 31d01318..809048b5 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup.go @@ -170,7 +170,8 @@ func maybeParseQuickStartArgs(args *types.CliArgs) (*types.ProxyInstanceConfig, // quick start instances don't have a way to configure otel, so just disable it otel := &types.OtelConfig{ - Enabled: false, + Enabled: false, + ServiceVersion: args.ProtocolVersion.String(), } instanceConfig := NewProxyInstanceConfig(args, args.QuickStartPort, otel, bigtableConfig) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup_test.go index ca90d065..6cc8213f 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/setup_test.go @@ -220,7 +220,8 @@ func TestParseProxyConfig(t *testing.T) { DefaultIntRowKeyEncoding: types.OrderedCodeEncoding, }, OtelConfig: &types.OtelConfig{ - Enabled: false, + Enabled: false, + ServiceVersion: "ProtocolVersion OSS 4", }, }, }, @@ -259,13 +260,14 @@ func TestParseProxyConfig(t *testing.T) { }, SchemaMappingTable: "schema_mapping", Session: &types.Session{ - GrpcChannels: 1, + GrpcChannels: 4, }, DefaultColumnFamily: "cf1", DefaultIntRowKeyEncoding: types.OrderedCodeEncoding, }, OtelConfig: &types.OtelConfig{ - Enabled: false, + Enabled: false, + ServiceVersion: "ProtocolVersion OSS 4", }, }, }, @@ -292,13 +294,14 @@ func TestParseProxyConfig(t *testing.T) { }, SchemaMappingTable: "sm", Session: &types.Session{ - GrpcChannels: 1, + GrpcChannels: 4, }, DefaultColumnFamily: "df", DefaultIntRowKeyEncoding: types.OrderedCodeEncoding, }, OtelConfig: &types.OtelConfig{ - Enabled: false, + Enabled: false, + ServiceVersion: "ProtocolVersion OSS 4", }, }, }, @@ -325,13 +328,14 @@ func TestParseProxyConfig(t *testing.T) { }, SchemaMappingTable: "sm", Session: &types.Session{ - GrpcChannels: 1, + GrpcChannels: 4, }, DefaultColumnFamily: "df", DefaultIntRowKeyEncoding: types.OrderedCodeEncoding, }, OtelConfig: &types.OtelConfig{ - Enabled: false, + Enabled: false, + ServiceVersion: "ProtocolVersion OSS 4", }, }, }, diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/types.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/types.go index eaae6759..c6cf63aa 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/types.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/types.go @@ -28,10 +28,12 @@ type yamlOtelConfig struct { Endpoint string `yaml:"endpoint"` } `yaml:"healthcheck"` Metrics struct { - Endpoint string `yaml:"endpoint"` + Endpoint string `yaml:"endpoint"` + GcpMetricsEnabled bool `yaml:"gcpMetricsEnabled"` } `yaml:"metrics"` Traces struct { Endpoint string `yaml:"endpoint"` + ProjectId string `yaml:"projectId"` SamplingRatio float64 `yaml:"samplingRatio"` } `yaml:"traces"` } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/utils.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/utils.go index 3ab04b6c..0b742576 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/utils.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/config/utils.go @@ -41,8 +41,9 @@ func loadProxyConfigFile(config *yamlProxyConfig, args *types.CliArgs) ([]*types // todo unit test with parent configs not defined to ensure no NPE is thrown otel := &types.OtelConfig{ - Enabled: config.Otel.Enabled, - ServiceName: config.Otel.ServiceName, + Enabled: config.Otel.Enabled, + ServiceName: config.Otel.ServiceName, + ServiceVersion: args.ProtocolVersion.String(), HealthCheck: struct { Enabled bool Endpoint string @@ -51,15 +52,19 @@ func loadProxyConfigFile(config *yamlProxyConfig, args *types.CliArgs) ([]*types Endpoint: config.Otel.HealthCheck.Endpoint, }, Metrics: struct { - Endpoint string + Endpoint string + GcpMetricsEnabled bool }{ - Endpoint: config.Otel.Metrics.Endpoint, + Endpoint: config.Otel.Metrics.Endpoint, + GcpMetricsEnabled: config.Otel.Metrics.GcpMetricsEnabled, }, Traces: struct { + ProjectId string Endpoint string SamplingRatio float64 }{ Endpoint: config.Otel.Traces.Endpoint, + ProjectId: config.Otel.Traces.ProjectId, SamplingRatio: config.Otel.Traces.SamplingRatio, }, } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy.go index c4eed9a3..fdf4c2c3 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy.go @@ -15,21 +15,16 @@ package proxy import ( - "bytes" "context" - "crypto/md5" "errors" "fmt" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors" - "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser" - "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/responsehandler" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/request_handlers" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/system_tables" - cql "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/cqlparser" - "io" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/trace" "net" - "strings" "sync" - "time" bigtableModule "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" @@ -37,11 +32,9 @@ import ( otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators" - "github.com/datastax/go-cassandra-native-protocol/frame" "github.com/datastax/go-cassandra-native-protocol/message" "github.com/datastax/go-cassandra-native-protocol/primitive" lru "github.com/hashicorp/golang-lru" - "go.opentelemetry.io/otel/attribute" "go.uber.org/zap" ) @@ -50,32 +43,20 @@ var ErrProxyAlreadyConnected = errors.New("proxy already connected") var ErrProxyNotConnected = errors.New("proxy not connected") const preparedIdSize = 16 -const Query = "CqlQuery" const translatorErrorMessage = "Error occurred at translators" const errorAtBigtable = "Error occurred at bigtable - " const errorWhileDecoding = "Error while decoding bytes - " const unhandledScenario = "Unhandled execution Scenario for prepared CqlQuery" -const errQueryNotPrepared = "query is not prepared" -const ( - handleQuery = "handleQuery" - handleBatch = "Batch" -) -// Events const ( - executingBigtableRequestEvent = "Executing Bigtable Mutation Request" - executingBigtableSQLAPIRequestEvent = "Executing Bigtable SQL API Request" - bigtableExecutionDoneEvent = "bigtable Execution Done" - gotBulkApplyResp = "Got the response for bulk apply" - sendingBulkApplyMutation = "Sending Mutation For Bulk Apply" + traceNamespace = "cassandra.bigtable.proxy" + handleOptions = traceNamespace + "/Options" ) -type Proxy struct { +type Server struct { ctx context.Context config *types.ProxyInstanceConfig logger *zap.Logger - cluster *proxycore.Cluster - sessions [primitive.ProtocolVersionDse2 + 1]sync.Map // Cache sessions per protocol version mu sync.Mutex isConnected bool isClosing bool @@ -83,7 +64,6 @@ type Proxy struct { listeners map[*net.Listener]struct{} eventClients sync.Map preparedQueryCache proxycore.PreparedCache[types.IPreparedQuery] - queryMetadataCache proxycore.PreparedCache[*message.PreparedResult] systemLocalValues map[string]message.Column closed chan struct{} localNode *node @@ -95,7 +75,47 @@ type Proxy struct { translator *translators.TranslatorManager executor *executors.QueryExecutorManager otelInst *otelgo.OpenTelemetry + tracer trace.Tracer otelShutdown func(context.Context) error + handlerManager *request_handlers.HandlerManager +} + +func (p *Server) CQLVersion() string { + return p.config.Options.CQLVersion +} + +// HandlePostDDLEvent handles common operations after DDL statements (CREATE, ALTER, DROP) +func (p *Server) HandlePostDDLEvent(queryType types.QueryType, keyspace types.Keyspace, table types.TableName) { + p.logger.Debug("sending post DDL event", zap.String("queryType", queryType.String()), zap.String("keyspace", string(keyspace)), zap.String("table", string(table))) + + var changeType primitive.SchemaChangeType + switch queryType { + case types.QueryTypeCreate: + changeType = primitive.SchemaChangeTypeCreated + case types.QueryTypeAlter: + changeType = primitive.SchemaChangeTypeUpdated + case types.QueryTypeDrop: + changeType = primitive.SchemaChangeTypeDropped + default: + p.logger.Warn("unhandled ddl event type", zap.String("queryType", queryType.String())) + return + } + + // SendEvent all clients of schema change + event := &proxycore.SchemaChangeEvent{ + Message: &message.SchemaChangeEvent{ + ChangeType: changeType, + Target: primitive.SchemaChangeTargetTable, + Keyspace: string(keyspace), + Object: string(table), + }, + } + p.eventClients.Range(func(key, _ interface{}) bool { + if client, ok := key.(*client); ok { + client.handleEvent(event) + } + return true + }) } type node struct { @@ -104,14 +124,14 @@ type node struct { tokens []string } -func (p *Proxy) OnEvent(event proxycore.Event) { +func (p *Server) OnEvent(event proxycore.Event) { switch evt := event.(type) { case *proxycore.SchemaChangeEvent: p.logger.Debug("Schema change event detected", zap.String("SchemaChangeEvent", evt.Message.String())) } } -func NewProxy(ctx context.Context, logger *zap.Logger, config *types.ProxyInstanceConfig) (*Proxy, error) { +func NewProxy(ctx context.Context, logger *zap.Logger, config *types.ProxyInstanceConfig) (*Server, error) { clientManager, err := types.CreateBigtableClientManager(ctx, config) if err != nil { return nil, err @@ -123,39 +143,23 @@ func NewProxy(ctx context.Context, logger *zap.Logger, config *types.ProxyInstan return nil, err } - bigtableClient := bigtableModule.NewBigtableClient(clientManager, logger, config.BigtableConfig, metadataStore) - - translator := translators.NewTranslatorManager(logger, metadataStore.Schemas(), config.BigtableConfig) - - // Enable OpenTelemetry traces by setting environment variable GOOGLE_API_GO_EXPERIMENTAL_TELEMETRY_PLATFORM_TRACING to the case-insensitive value "opentelemetry" before loading the client library. - otelConfig := &otelgo.OTelConfig{} - otelInst := &otelgo.OpenTelemetry{Config: &otelgo.OTelConfig{OTELEnabled: false}} - var shutdownOTel func(context.Context) error - // Initialize OpenTelemetry - if config.OtelConfig.Enabled { - otelConfig = &otelgo.OTelConfig{ - TracerEndpoint: config.OtelConfig.Traces.Endpoint, - MetricEndpoint: config.OtelConfig.Metrics.Endpoint, - ServiceName: config.OtelConfig.ServiceName, - OTELEnabled: config.OtelConfig.Enabled, - TraceSampleRatio: config.OtelConfig.Traces.SamplingRatio, - HealthCheckEnabled: config.OtelConfig.HealthCheck.Enabled, - HealthCheckEp: config.OtelConfig.HealthCheck.Endpoint, - ServiceVersion: config.Options.ProtocolVersion.String(), - } - } else { - otelConfig = &otelgo.OTelConfig{OTELEnabled: false} - } - otelInst, shutdownOTel, err = otelgo.NewOpenTelemetry(ctx, otelConfig, logger) + var otelInst *otelgo.OpenTelemetry + otelInst, shutdownOTel, err = otelgo.NewOpenTelemetry(ctx, config.OtelConfig, logger) if err != nil { logger.Error("Failed to enable the OTEL: " + err.Error()) return nil, err } + bigtableClient := bigtableModule.NewBigtableClient(clientManager, logger, config.BigtableConfig, metadataStore) + + translator := translators.NewTranslatorManager(logger, metadataStore.Schemas(), config.BigtableConfig) + systemTables := system_tables.NewSystemTableManager(metadataStore, logger) - proxy := &Proxy{ + handlers := request_handlers.NewHandlerManager() + + proxy := &Server{ ctx: ctx, config: config, logger: logger, @@ -167,11 +171,15 @@ func NewProxy(ctx context.Context, logger *zap.Logger, config *types.ProxyInstan metadataStore: metadataStore, bigtableClient: bigtableClient, translator: translator, - executor: executors.NewQueryExecutorManager(logger, metadataStore.Schemas(), bigtableClient, systemTables.Db()), + executor: executors.NewQueryExecutorManager(logger, metadataStore.Schemas(), bigtableClient, systemTables.Db(), otelInst), otelInst: otelInst, + tracer: otel.GetTracerProvider().Tracer("handler"), otelShutdown: shutdownOTel, + handlerManager: handlers, } + handlers.InitHandlers(proxy) + err = systemTables.Initialize(proxy) if err != nil { logger.Error("Failed to initialize system table manager: " + err.Error()) @@ -181,7 +189,7 @@ func NewProxy(ctx context.Context, logger *zap.Logger, config *types.ProxyInstan return proxy, nil } -func (p *Proxy) Connect() error { +func (p *Server) Connect() error { p.mu.Lock() defer p.mu.Unlock() @@ -195,38 +203,12 @@ func (p *Proxy) Connect() error { if err != nil { return fmt.Errorf("unable to create cache: %w", err) } - p.queryMetadataCache, err = NewDefaultPreparedCache[*message.PreparedResult](cacheSize) - if err != nil { - return fmt.Errorf("unable to create cache: %w", err) - } - - // connecting to cassandra cluster - p.cluster, err = proxycore.ConnectCluster(p.ctx, proxycore.ClusterConfig{ - Version: p.config.Options.ProtocolVersion, - Logger: p.logger, - }) - - if err != nil { - return fmt.Errorf("unable to connect to cluster %w", err) - } err = p.buildNodes() if err != nil { return fmt.Errorf("unable to build node information: %w", err) } - // Create cassandra session - sess, err := proxycore.ConnectSession(p.ctx, proxycore.SessionConfig{ - Version: p.cluster.NegotiatedVersion, - Logger: p.logger, - }) - - if err != nil { - return fmt.Errorf("unable to connect session: %w", err) - } - - p.sessions[p.cluster.NegotiatedVersion].Store("", sess) - err = p.systemTableManager.ReloadSystemTables() if err != nil { p.logger.Error("failed to update system tables", zap.Error(err)) @@ -238,7 +220,7 @@ func (p *Proxy) Connect() error { // Serve the proxy using the specified listener. It can be called multiple times with different listeners allowing // them to share the same backend clusters. -func (p *Proxy) Serve(l net.Listener) (err error) { +func (p *Server) Serve(l net.Listener) (err error) { l = &closeOnceListener{Listener: l} defer l.Close() @@ -261,7 +243,7 @@ func (p *Proxy) Serve(l net.Listener) (err error) { } } -func (p *Proxy) GetSystemTableConfig() system_tables.SystemTableConfig { +func (p *Server) GetSystemTableConfig() system_tables.SystemTableConfig { var peers []system_tables.PeerConfig for _, n := range p.nodes { peers = append(peers, system_tables.PeerConfig{ @@ -283,7 +265,7 @@ func (p *Proxy) GetSystemTableConfig() system_tables.SystemTableConfig { return systemTableConfig } -func (p *Proxy) addListener(l *net.Listener) error { +func (p *Server) addListener(l *net.Listener) error { p.mu.Lock() defer p.mu.Unlock() if p.isClosing { @@ -296,13 +278,49 @@ func (p *Proxy) addListener(l *net.Listener) error { return nil } -func (p *Proxy) removeListener(l *net.Listener) { +func (p *Server) removeListener(l *net.Listener) { p.mu.Lock() defer p.mu.Unlock() delete(p.listeners, l) } -func (p *Proxy) Close() error { +func (p *Server) Config() *types.ProxyInstanceConfig { + return p.config +} + +func (p *Server) OtelInst() *otelgo.OpenTelemetry { + return p.otelInst +} + +func (p *Server) Logger() *zap.Logger { + return p.logger +} + +func (p *Server) Translator() *translators.TranslatorManager { + return p.translator +} + +func (p *Server) Executor() *executors.QueryExecutorManager { + return p.executor +} + +func (p *Server) PreparedQueryCache() proxycore.PreparedCache[types.IPreparedQuery] { + return p.preparedQueryCache +} + +func (p *Server) BigtableClient() *bigtableModule.BigtableAdapter { + return p.bigtableClient +} + +func (p *Server) EventClients() *sync.Map { + return &p.eventClients +} + +func (p *Server) HandleOptions() string { + return handleOptions +} + +func (p *Server) Close() error { p.mu.Lock() defer p.mu.Unlock() select { @@ -329,7 +347,7 @@ func (p *Proxy) Close() error { return err } -func (p *Proxy) Ready() bool { +func (p *Server) Ready() bool { return true } @@ -342,7 +360,7 @@ func getUdsPeerCredentials(conn *net.UnixConn) (UCred, error) { return getUdsPeerCredentialsOS(conn) } -func (p *Proxy) handle(conn net.Conn) { +func (p *Server) handle(conn net.Conn) { if tcpConn, ok := conn.(*net.TCPConn); ok { if err := tcpConn.SetKeepAlive(false); err != nil { p.logger.Warn("failed to disable keepalive on connection", zap.Error(err)) @@ -361,7 +379,6 @@ func (p *Proxy) handle(conn net.Conn) { } cl := &client{ - ctx: p.ctx, proxy: p, } cl.sender = cl @@ -370,32 +387,24 @@ func (p *Proxy) handle(conn net.Conn) { cl.conn.Start() } -func (p *Proxy) buildNodes() (err error) { - localDC := p.config.DC - if len(localDC) == 0 { - localDC = p.cluster.Info.LocalDC - p.logger.Info("no local DC configured using DC from the first successful contact point", - zap.String("dc", localDC)) - } - +func (p *Server) buildNodes() (err error) { p.localNode = &node{ - dc: localDC, + dc: p.config.DC, } - return nil } -func (p *Proxy) addClient(cl *client) { +func (p *Server) addClient(cl *client) { p.mu.Lock() defer p.mu.Unlock() p.clients[cl] = struct{}{} } -func (p *Proxy) registerForEvents(cl *client) { +func (p *Server) registerForEvents(cl *client) { p.eventClients.Store(cl, struct{}{}) } -func (p *Proxy) removeClient(cl *client) { +func (p *Server) removeClient(cl *client) { p.eventClients.Delete(cl) p.mu.Lock() @@ -404,341 +413,6 @@ func (p *Proxy) removeClient(cl *client) { } -type Sender interface { - Send(hdr *frame.Header, msg message.Message) -} - -type client struct { - sessionKeyspace types.Keyspace - ctx context.Context - proxy *Proxy - conn *proxycore.Conn - sender Sender -} - -func (c *client) SetSessionKeyspace(k types.Keyspace) { - c.sessionKeyspace = k -} - -func (c *client) Receive(reader io.Reader) error { - raw, err := codec.DecodeRawFrame(reader) - if err != nil { - if !errors.Is(err, io.EOF) { - c.proxy.logger.Error("unable to decode frame", zap.Error(err)) - } - return err - } - - if raw.Header.Version > c.proxy.config.Options.MaxProtocolVersion || raw.Header.Version < primitive.ProtocolVersion3 { - c.sender.Send(raw.Header, &message.ProtocolError{ - ErrorMessage: fmt.Sprintf("Invalid or unsupported protocol version %d", raw.Header.Version), - }) - return nil - } - - body, err := codec.DecodeBody(raw.Header, bytes.NewReader(raw.Body)) - if err != nil { - c.proxy.logger.Error("unable to decode body", zap.Error(err)) - return err - } - - switch msg := body.Message.(type) { - case *message.Options: - // CC - responding with status READY - c.sender.Send(raw.Header, &message.Supported{Options: map[string][]string{ - "CQL_VERSION": {c.proxy.config.Options.CQLVersion}, - "COMPRESSION": {}, - }}) - case *message.Startup: - // CC - register for Event types and respond READY - c.sender.Send(raw.Header, &message.Ready{}) - case *message.Register: - c.proxy.logger.Info("Client registered for events", zap.Any("event_types", msg.EventTypes)) - for _, t := range msg.EventTypes { - if t == primitive.EventTypeSchemaChange { - c.proxy.registerForEvents(c) - } - } - c.sender.Send(raw.Header, &message.Ready{}) - case *message.Prepare: - c.handlePrepare(raw, msg) - case *partialExecute: - c.handleExecute(raw, msg) - case *partialQuery: - c.handleQuery(raw, msg) - case *partialBatch: - c.handleBatch(raw, msg) - default: - c.sender.Send(raw.Header, &message.ProtocolError{ErrorMessage: "Unsupported operation"}) - } - return nil -} - -func (c *client) handlePrepare(raw *frame.RawFrame, msg *message.Prepare) { - id := c.getQueryId(msg) - if response, found := c.proxy.queryMetadataCache.Load(id); found { - c.sender.Send(raw.Header, response) - return - } - - c.proxy.logger.Debug("preparing query", zap.String(Query, msg.Query), zap.Int16("stream", raw.Header.StreamId)) - - keyspace := c.sessionKeyspace - if len(msg.Keyspace) != 0 { - keyspace = types.Keyspace(msg.Keyspace) - } - - p := parser.NewParser(msg.Query) - qt, err := parseQueryType(p) - if err != nil { - c.proxy.logger.Error("error parsing query to see if it's handled", zap.String(Query, msg.Query), zap.Error(err)) - c.sender.Send(raw.Header, &message.Invalid{ErrorMessage: err.Error()}) - return - } - rawQuery := types.NewRawQuery(raw.Header, keyspace, msg.Query, p, qt) - - c.handleServerPreparedQuery(rawQuery, id) -} - -func parseQueryType(p *parser.ProxyCqlParser) (types.QueryType, error) { - tok := p.GetFirstToken() - t := tok.GetTokenType() - switch t { - case cql.CqlLexerK_SELECT: - return types.QueryTypeSelect, nil - case cql.CqlLexerK_INSERT: - return types.QueryTypeInsert, nil - case cql.CqlLexerK_UPDATE: - return types.QueryTypeUpdate, nil - case cql.CqlLexerK_DELETE: - return types.QueryTypeDelete, nil - case cql.CqlLexerK_CREATE: - return types.QueryTypeCreate, nil - case cql.CqlLexerK_ALTER: - return types.QueryTypeAlter, nil - case cql.CqlLexerK_DROP: - return types.QueryTypeDrop, nil - case cql.CqlLexerK_TRUNCATE: - return types.QueryTypeTruncate, nil - case cql.CqlLexerK_USE: - return types.QueryTypeUse, nil - case cql.CqlLexerK_DESCRIBE, cql.CqlLexerK_DESC: - return types.QueryTypeDescribe, nil - default: - return types.QueryTypeUnknown, fmt.Errorf("unsupported query type: %s", tok.String()) - } -} - -func (c *client) getQueryId(msg *message.Prepare) [16]byte { - // Generating unique prepared query_id - return md5.Sum([]byte(msg.Query + string(c.sessionKeyspace))) -} - -// handleServerPreparedQuery handle prepared query that was supposed to run on cassandra server -// This method will keep track of prepared query in a map and send hashed query_id with result -// metadata and variable column metadata to the client -// -// Parameters: -// - raw: *frame.RawFrame -// - msg: *message.Prepare -// -// Returns: nil -func (c *client) handleServerPreparedQuery(query *types.RawQuery, id [16]byte) { - preparedQuery, err := c.prepareQuery(query) - if err != nil { - c.proxy.logger.Error(translatorErrorMessage, zap.String(Query, query.RawCql()), zap.Error(err)) - c.sender.Send(query.Header(), &message.Invalid{ErrorMessage: err.Error()}) - return - } - - response, err := responsehandler.BuildPreparedResultResponse(id, preparedQuery) - - // update caches - c.proxy.preparedQueryCache.Store(id, preparedQuery) - c.proxy.queryMetadataCache.Store(id, response) - - c.sender.Send(query.Header(), response) -} - -func (c *client) prepareQuery(query *types.RawQuery) (types.IPreparedQuery, error) { - preparedQuery, err := c.proxy.translator.TranslateQuery(query, c.sessionKeyspace) - if err != nil { - return nil, err - } - - btPreparedQuery, err := c.proxy.bigtableClient.PrepareStatement(c.ctx, preparedQuery) - if err != nil { - return nil, fmt.Errorf("failed to prepare bigtable statement `%s`: %w", preparedQuery.BigtableQuery(), err) - } - preparedQuery.SetBigtablePreparedQuery(btPreparedQuery) - - return preparedQuery, err -} - -// handleExecute for prepared query -func (c *client) handleExecute(raw *frame.RawFrame, msg *partialExecute) { - ctx := context.Background() - id := preparedIdKey(msg.queryId) - - preparedStmt, ok := c.proxy.preparedQueryCache.Load(id) - if !ok { - c.proxy.logger.Error(errQueryNotPrepared) - c.sender.Send(raw.Header, &message.ServerError{ErrorMessage: errQueryNotPrepared}) - return - } - - boundQuery, err := c.proxy.translator.BindQuery(preparedStmt, msg.PositionalValues, msg.NamedValues, raw.Header.Version) - if err != nil { - c.proxy.logger.Error(errorWhileDecoding, zap.String(Query, preparedStmt.CqlQuery()), zap.Error(err)) - c.sender.Send(raw.Header, &message.ConfigError{ErrorMessage: err.Error()}) - return - } - - results, err := c.proxy.executor.Execute(ctx, c, boundQuery) - if err != nil { - c.proxy.logger.Error(errorAtBigtable, zap.String(Query, preparedStmt.CqlQuery()), zap.String(Query, preparedStmt.BigtableQuery()), zap.Error(err)) - c.sender.Send(raw.Header, &message.ConfigError{ErrorMessage: err.Error()}) - return - } - if preparedStmt.QueryType().IsDDLType() { - c.handlePostDDLEvent(preparedStmt.QueryType(), preparedStmt.Keyspace(), preparedStmt.Table()) - } - c.sender.Send(raw.Header, results) -} - -// handle batch queries -func (c *client) handleBatch(raw *frame.RawFrame, msg *partialBatch) { - startTime := time.Now() - otelCtx, span := c.proxy.otelInst.StartSpan(c.proxy.ctx, handleBatch, []attribute.KeyValue{ - attribute.Int("Batch Size", len(msg.queryOrIds)), - }) - defer c.proxy.otelInst.EndSpan(span) - var otelErr error - defer c.proxy.otelInst.RecordMetrics(otelCtx, handleBatch, startTime, handleBatch, c.sessionKeyspace, otelErr) - bulkMutations, keyspace, err := c.bindBulkOperations(msg, raw.Header.Version) - if err != nil { - c.proxy.logger.Error("Error preparing batch query metadata", zap.Error(err)) - c.sender.Send(raw.Header, &message.ConfigError{ErrorMessage: err.Error()}) - otelErr = err - c.proxy.otelInst.RecordError(span, otelErr) - return - } - otelgo.AddAnnotation(otelCtx, sendingBulkApplyMutation) - var errs []string - for tableName, mutations := range bulkMutations.Mutations() { - res, err := c.proxy.bigtableClient.ApplyBulkMutation(otelCtx, keyspace, tableName, mutations) - if err != nil || res.FailedRows != "" { - c.proxy.otelInst.RecordError(span, err) - errs = append(errs, res.FailedRows) - } - } - otelgo.AddAnnotation(otelCtx, gotBulkApplyResp) - if len(errs) > 0 { - c.sender.Send(raw.Header, &message.ServerError{ErrorMessage: strings.Join(errs, "\n")}) - return - } - c.sender.Send(raw.Header, &message.VoidResult{}) -} - -func (c *client) bindBulkOperations(msg *partialBatch, pv primitive.ProtocolVersion) (*bigtableModule.BigtableBulkMutation, types.Keyspace, error) { - var keyspace types.Keyspace - tableMutationsMap := bigtableModule.NewBigtableBulkMutation() - for index, queryId := range msg.queryOrIds { - queryOrId, ok := queryId.([]byte) - if !ok { - return nil, "", fmt.Errorf("batch query id malformed") - } - id := preparedIdKey(queryOrId) - preparedStmt, ok := c.proxy.preparedQueryCache.Load(id) - if !ok { - return nil, "", fmt.Errorf("prepared query not found in cache") - } - - if preparedStmt.Keyspace() != "" { - keyspace = preparedStmt.Keyspace() - } - - // note: we don't support batch named queries at this time - executableQuery, err := c.proxy.translator.BindQuery(preparedStmt, msg.BatchPositionalValues[index], nil, pv) - if err != nil { - return nil, "", err - } - mutation, ok := executableQuery.AsBulkMutation() - if !ok { - return nil, "", fmt.Errorf("query type '%s' not compatible with bulk", executableQuery.QueryType().String()) - } - tableMutationsMap.AddMutation(mutation) - } - if keyspace == "" { - keyspace = c.sessionKeyspace - } - return tableMutationsMap, keyspace, nil -} - -func (c *client) handleQuery(raw *frame.RawFrame, msg *partialQuery) { - startTime := time.Now() - c.proxy.logger.Debug("handling query", zap.String("encodedQuery", msg.query), zap.Int16("stream", raw.Header.StreamId)) - otelCtx, span := c.proxy.otelInst.StartSpan(c.proxy.ctx, handleQuery, []attribute.KeyValue{ - attribute.String("CqlQuery", msg.query), - }) - - p := parser.NewParser(msg.query) - qt, err := parseQueryType(p) - if err != nil { - c.proxy.logger.Error("error parsing query to see if it's handled", zap.String(Query, msg.query), zap.Error(err)) - c.sender.Send(raw.Header, &message.Invalid{ErrorMessage: err.Error()}) - return - } - - rawQuery := types.NewRawQuery(raw.Header, c.sessionKeyspace, msg.query, p, qt) - defer c.proxy.otelInst.EndSpan(span) - - var otelErr error - defer c.proxy.otelInst.RecordMetrics(otelCtx, handleQuery, startTime, rawQuery.QueryType().String(), c.sessionKeyspace, otelErr) - - query, err := c.prepareQuery(rawQuery) - if err != nil { - c.proxy.logger.Error(translatorErrorMessage, zap.String(Query, msg.query), zap.Error(err)) - c.sender.Send(raw.Header, &message.Invalid{ErrorMessage: err.Error()}) - return - } - values := types.NewQueryParameterValues(query.Parameters(), time.Now()) - executableQuery, err := c.proxy.translator.BindQueryParameters(query, values, raw.Header.Version) - if err != nil { - c.proxy.logger.Error(translatorErrorMessage, zap.String(Query, msg.query), zap.Error(err)) - c.sender.Send(raw.Header, &message.Invalid{ErrorMessage: err.Error()}) - return - } - - otelgo.AddAnnotation(otelCtx, executingBigtableSQLAPIRequestEvent) - selectResult, err := c.proxy.executor.Execute(otelCtx, c, executableQuery) - otelgo.AddAnnotation(otelCtx, bigtableExecutionDoneEvent) - - if err != nil { - c.proxy.logger.Error(errorAtBigtable, zap.String(Query, msg.query), zap.Error(err)) - otelErr = err - c.proxy.otelInst.RecordError(span, otelErr) - c.sender.Send(raw.Header, &message.Invalid{ErrorMessage: err.Error()}) - return - } - - if rawQuery.QueryType().IsDDLType() { - c.handlePostDDLEvent(query.QueryType(), query.Keyspace(), query.Table()) - } - - c.sender.Send(raw.Header, selectResult) -} - -func (c *client) Send(hdr *frame.Header, msg message.Message) { - _ = c.conn.Write(proxycore.SenderFunc(func(writer io.Writer) error { - return codec.EncodeFrame(frame.NewFrame(hdr.Version, hdr.StreamId, msg), writer) - })) -} - -func (c *client) Closing(_ error) { - c.proxy.removeClient(c) -} - // NewDefaultPreparedCache creates a new default prepared cache capping the max item capacity to `size`. func NewDefaultPreparedCache[T any](size int) (proxycore.PreparedCache[T], error) { cache, err := lru.New(size) @@ -782,48 +456,3 @@ func (oc *closeOnceListener) Close() error { } func (oc *closeOnceListener) close() { oc.closeErr = oc.Listener.Close() } - -// handleEvent handles events from the proxy core -// It sends the event message to all connected clients. -func (c *client) handleEvent(event proxycore.Event) { - switch evt := event.(type) { - case *proxycore.SchemaChangeEvent: - c.sender.Send(&frame.Header{ - Version: c.proxy.config.Options.ProtocolVersion, - StreamId: -1, // -1 for events - OpCode: primitive.OpCodeEvent, - }, evt.Message) - } -} - -// handlePostDDLEvent handles common operations after DDL statements (CREATE, ALTER, DROP) -func (c *client) handlePostDDLEvent(queryType types.QueryType, keyspace types.Keyspace, table types.TableName) { - var changeType primitive.SchemaChangeType - switch queryType { - case types.QueryTypeCreate: - changeType = primitive.SchemaChangeTypeCreated - case types.QueryTypeAlter: - changeType = primitive.SchemaChangeTypeUpdated - case types.QueryTypeDrop: - changeType = primitive.SchemaChangeTypeDropped - default: - c.proxy.logger.Warn("unhandled ddl event type", zap.String("queryType", queryType.String())) - return - } - - // SendEvent all clients of schema change - event := &proxycore.SchemaChangeEvent{ - Message: &message.SchemaChangeEvent{ - ChangeType: changeType, - Target: primitive.SchemaChangeTargetTable, - Keyspace: string(keyspace), - Object: string(table), - }, - } - c.proxy.eventClients.Range(func(key, _ interface{}) bool { - if client, ok := key.(*client); ok { - client.handleEvent(event) - } - return true - }) -} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types/request_handler.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types/request_handler.go new file mode 100644 index 00000000..248bf8f7 --- /dev/null +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/proxy_types/request_handler.go @@ -0,0 +1,123 @@ +package proxy_types + +import ( + "encoding/hex" + bigtableModule "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/bigtable" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/executors" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" + otelgo "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/otel" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore" + "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators" + "github.com/datastax/go-cassandra-native-protocol/message" + "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.uber.org/zap" + "sync" +) + +// Constants moved from proxy package +const ( + Query = "CqlQuery" + QueryTypeConst = "QueryType" // Renamed from QueryType to avoid collision with types.QueryType + ErrQueryNotPrepared = "query is not prepared" +) + +// Trace names +const ( + TraceNamespace = "cassandra.bigtable.proxy" + HandleQuery = TraceNamespace + "/HandleQuery" + HandleBatch = TraceNamespace + "/ExecuteBatch" + HandlePrepare = TraceNamespace + "/PrepareQuery" + HandleExecute = TraceNamespace + "/ExecuteQuery" + HandleRegister = TraceNamespace + "/Register" + HandleStartup = TraceNamespace + "/Startup" + HandleOptions = TraceNamespace + "/Options" +) + +// Events +const ( + ExecutingBigtableRequestEvent = "Executing Bigtable Mutation Request" + ExecutingBigtableSQLAPIRequestEvent = "Executing Bigtable SQL API Request" + BigtableExecutionDoneEvent = "bigtable Execution Done" + GotBulkApplyResp = "Got the response for bulk apply" + SendingBulkApplyMutation = "Sending Mutation For Bulk Apply" +) + +type IProxy interface { + Config() *types.ProxyInstanceConfig + OtelInst() *otelgo.OpenTelemetry + Logger() *zap.Logger + Translator() *translators.TranslatorManager + Executor() *executors.QueryExecutorManager + PreparedQueryCache() proxycore.PreparedCache[types.IPreparedQuery] + BigtableClient() *bigtableModule.BigtableAdapter + EventClients() *sync.Map +} + +func PreparedIdKey(bytes []byte) [16]byte { + var buf [16]byte + copy(buf[:], bytes) + return buf +} + +// Partial types moved from codecs.go + +type PartialQuery struct { + Query string +} + +func (p *PartialQuery) IsResponse() bool { + return false +} + +func (p *PartialQuery) GetOpCode() primitive.OpCode { + return primitive.OpCodeQuery +} + +func (p *PartialQuery) Clone() message.Message { + return &PartialQuery{p.Query} +} + +type PartialExecute struct { + QueryId []byte + PositionalValues []*primitive.Value + NamedValues map[string]*primitive.Value +} + +func (m *PartialExecute) IsResponse() bool { + return false +} + +func (m *PartialExecute) GetOpCode() primitive.OpCode { + return primitive.OpCodeExecute +} + +func (m *PartialExecute) Clone() message.Message { + return &PartialExecute{ + QueryId: primitive.CloneByteSlice(m.QueryId), + } +} + +func (m *PartialExecute) String() string { + return "EXECUTE " + hex.EncodeToString(m.QueryId) +} + +type PartialBatch struct { + QueryOrIds []interface{} + BatchPositionalValues [][]*primitive.Value +} + +func (p *PartialBatch) IsResponse() bool { + return false +} + +func (p *PartialBatch) GetOpCode() primitive.OpCode { + return primitive.OpCodeBatch +} + +func (p *PartialBatch) Clone() message.Message { + queryOrIds := make([]interface{}, len(p.QueryOrIds)) + copy(queryOrIds, p.QueryOrIds) + positionalValues := make([][]*primitive.Value, len(p.BatchPositionalValues)) + copy(positionalValues, p.BatchPositionalValues) + return &PartialBatch{queryOrIds, positionalValues} +} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/run.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/run.go index 66e5cea6..cb2e8d31 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/run.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxy/run.go @@ -70,7 +70,7 @@ func Run(ctx context.Context, args []string) error { } var mux http.ServeMux wg.Add(1) - go func(cfg *types.ProxyInstanceConfig, p *Proxy, mux *http.ServeMux) { + go func(cfg *types.ProxyInstanceConfig, p *Server, mux *http.ServeMux) { defer wg.Done() err := listenAndServe(listenerConfig, p, mux, ctx, logger) // Use cfg2 or other instances as needed if err != nil { @@ -86,7 +86,7 @@ func Run(ctx context.Context, args []string) error { } // listenAndServe correctly handles serving both the proxy and an HTTP server simultaneously. -func listenAndServe(c *types.ProxyInstanceConfig, p *Proxy, mux *http.ServeMux, ctx context.Context, logger *zap.Logger) (err error) { +func listenAndServe(c *types.ProxyInstanceConfig, p *Server, mux *http.ServeMux, ctx context.Context, logger *zap.Logger) (err error) { logger.Info("Starting proxy with configuration:\n") logger.Info(fmt.Sprintf(" Bind: %s\n", c.Bind)) logger.Info(fmt.Sprintf(" Use Unix Socket: %v\n", c.Options.UseUnixSocket)) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/cluster.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/cluster.go index 6ae2ddc3..8a93210d 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/cluster.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/cluster.go @@ -24,11 +24,6 @@ import ( "go.uber.org/zap" ) -const ( - DefaultRefreshWindow = 10 * time.Second - DefaultRefreshTimeout = 5 * time.Second -) - type Event interface { isEvent() // Marker method for the event interface } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/codecs.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/codecs.go index ec279b95..238c1ba1 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/codecs.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/codecs.go @@ -19,12 +19,9 @@ import ( "github.com/datastax/go-cassandra-native-protocol/datacodec" "github.com/datastax/go-cassandra-native-protocol/datatype" - "github.com/datastax/go-cassandra-native-protocol/frame" "github.com/datastax/go-cassandra-native-protocol/primitive" ) -var codec = frame.NewRawCodec() - var primitiveCodecs = map[datatype.DataType]datacodec.Codec{ datatype.Ascii: datacodec.Ascii, datatype.Bigint: datacodec.Bigint, diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/session.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/session.go index f1865699..0ba21f63 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/session.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/third_party/datastax/proxycore/session.go @@ -14,54 +14,8 @@ package proxycore -import ( - "context" - "sync" - "time" - - "github.com/datastax/go-cassandra-native-protocol/frame" - "github.com/datastax/go-cassandra-native-protocol/primitive" - "go.uber.org/zap" -) - -// PreparedEntry is an entry in the prepared cache. -type PreparedEntry struct { - PreparedFrame *frame.RawFrame -} - type PreparedCache[T any] interface { // Store add an entry to the cache. Store(id [16]byte, entry T) Load(id [16]byte) (entry T, ok bool) } - -type SessionConfig struct { - Version primitive.ProtocolVersion - Auth Authenticator - // PreparedCache a global cache share across sessions for storing previously prepared queries - ConnectTimeout time.Duration - HeartBeatInterval time.Duration - IdleTimeout time.Duration - Logger *zap.Logger -} - -type Session struct { - ctx context.Context - config SessionConfig - logger *zap.Logger - pools sync.Map - connected chan struct{} - failed chan error -} - -func ConnectSession(ctx context.Context, config SessionConfig) (*Session, error) { - session := &Session{ - ctx: ctx, - config: config, - logger: GetOrCreateNopLogger(config.Logger), - pools: sync.Map{}, - connected: make(chan struct{}), - failed: make(chan error, 1), - } - return session, nil -} diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/alter_translator/translator_alter_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/alter_translator/translator_alter_test.go index 74c0d5e8..fe90228e 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/alter_translator/translator_alter_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/alter_translator/translator_alter_test.go @@ -259,7 +259,7 @@ func TestTranslateAlterTableToBigtable(t *testing.T) { t.Run(tt.name, func(t *testing.T) { require.NotNil(t, tt.tableConfig, "tests must define a table config") tr := NewAlterTranslator(mockdata.GetSchemaMappingConfig()) - query := types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeAlter) + query := types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeAlter) got, err := tr.Translate(query, tt.defaultKeyspace) if tt.error != "" { require.Error(t, err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/create_translator/translator_create_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/create_translator/translator_create_test.go index 64e3e258..d0e74acf 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/create_translator/translator_create_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/create_translator/translator_create_test.go @@ -465,7 +465,7 @@ func TestTranslateCreateTableToBigtable(t *testing.T) { SchemaMappingTable: "schema_mapping", DefaultIntRowKeyEncoding: tt.defaultIntRowKeyEncoding, }) - query := types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeCreate) + query := types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeCreate) got, err := tr.Translate(query, tt.defaultKeyspace) if tt.error != "" { require.Error(t, err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/delete_translator/translator_delete_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/delete_translator/translator_delete_test.go index 9bff2dd8..006b8c8c 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/delete_translator/translator_delete_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/delete_translator/translator_delete_test.go @@ -187,7 +187,7 @@ func TestTranslator_TranslateDeleteQuerytoBigtable(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tr := NewDeleteTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) if tt.wantErr != "" { require.Error(t, err) assert.Contains(t, err.Error(), tt.wantErr) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/desc_translator/translator_describe_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/desc_translator/translator_describe_test.go index d2c7413c..ed5e0eee 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/desc_translator/translator_describe_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/desc_translator/translator_describe_test.go @@ -101,7 +101,7 @@ func TestTranslator_TranslateDescribeQuery(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tr := NewDescTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) if tt.wantErr != "" { require.Error(t, err) assert.Contains(t, err.Error(), tt.wantErr) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/drop_translator/translator_drop_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/drop_translator/translator_drop_test.go index 23695749..68f10fd7 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/drop_translator/translator_drop_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/drop_translator/translator_drop_test.go @@ -141,7 +141,7 @@ func TestTranslateDropTableToBigtable(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tr := NewDropTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) if tt.error != "" { require.Error(t, err) assert.Contains(t, err.Error(), tt.error) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/insert_translator/translator_insert_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/insert_translator/translator_insert_test.go index 17edb82d..ac5e6125 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/insert_translator/translator_insert_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/insert_translator/translator_insert_test.go @@ -378,7 +378,7 @@ func TestSelectTranslator_Translate(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tr := NewInsertTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeInsert), tt.sessionKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeInsert), tt.sessionKeyspace) if tt.wantErr != "" { require.Error(t, err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/manager.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/manager.go index 09a08987..cba7668c 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/manager.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/manager.go @@ -17,6 +17,7 @@ package translators import ( + "context" "fmt" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" schemaMapping "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata" @@ -32,6 +33,9 @@ import ( "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/update_translator" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/use_translator" "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" "go.uber.org/zap" ) @@ -40,12 +44,13 @@ type TranslatorManager struct { SchemaMappingConfig *schemaMapping.SchemaMetadata translators map[types.QueryType]types.IQueryTranslator config *types.BigtableConfig + tracer trace.Tracer } func NewTranslatorManager(logger *zap.Logger, schemaMappingConfig *schemaMapping.SchemaMetadata, config *types.BigtableConfig) *TranslatorManager { // add more translators here translators := []types.IQueryTranslator{ - select_translator.NewSelectTranslator(schemaMappingConfig), + select_translator.NewSelectTranslator(schemaMappingConfig, logger), insert_translator.NewInsertTranslator(schemaMappingConfig), update_translator.NewUpdateTranslator(schemaMappingConfig), delete_translator.NewDeleteTranslator(schemaMappingConfig), @@ -69,18 +74,25 @@ func NewTranslatorManager(logger *zap.Logger, schemaMappingConfig *schemaMapping SchemaMappingConfig: schemaMappingConfig, translators: tm, config: config, + tracer: otel.GetTracerProvider().Tracer("translator"), } } -func (t *TranslatorManager) TranslateQuery(q *types.RawQuery, sessionKeyspace types.Keyspace) (types.IPreparedQuery, error) { +func (t *TranslatorManager) TranslateQuery(ctx context.Context, q *types.RawQuery, sessionKeyspace types.Keyspace) (types.IPreparedQuery, error) { + _, span := t.tracer.Start(ctx, "translate") + defer span.End() + defer q.Parser().Release() + queryTranslator, err := t.getTranslator(q.QueryType()) if err != nil { + span.RecordError(err) return nil, err } preparedQuery, err := queryTranslator.Translate(q, sessionKeyspace) - if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) return nil, err } @@ -88,10 +100,13 @@ func (t *TranslatorManager) TranslateQuery(q *types.RawQuery, sessionKeyspace ty // ensure user doesn't try to drop or corrupt the schema mapping table if !preparedQuery.Keyspace().IsSystemKeyspace() && preparedQuery.Table() == t.config.SchemaMappingTable { - return nil, fmt.Errorf("table name cannot be the same as the configured schema mapping table name '%s'", t.config.SchemaMappingTable) + err = fmt.Errorf("table name cannot be the same as the configured schema mapping table name '%s'", t.config.SchemaMappingTable) + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return nil, err } - return preparedQuery, err + return preparedQuery, nil } func (t *TranslatorManager) BindQuery(st types.IPreparedQuery, cassandraValues []*primitive.Value, namedValues map[string]*primitive.Value, pv primitive.ProtocolVersion) (types.IExecutableQuery, error) { diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select.go index 216333e1..cf6ccd63 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select.go @@ -23,14 +23,18 @@ import ( "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/common" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/utilities" "github.com/datastax/go-cassandra-native-protocol/primitive" + "go.uber.org/zap" + "time" ) func (t *SelectTranslator) Translate(query *types.RawQuery, sessionKeyspace types.Keyspace) (types.IPreparedQuery, error) { + parserStart := time.Now() selectObj, err := query.Parser().Select_() if err != nil { return nil, err } - + parserDuration := time.Now().Sub(parserStart) + t.logger.Info("PARSER_TIMING", zap.String("duration", parserDuration.String()), zap.String("query", query.RawCql())) keyspaceName, tableName, err := common.ParseTableSpec(selectObj.FromSpec().TableSpec(), sessionKeyspace) if err != nil { return nil, err diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select_test.go index 5f830a3d..b29b8f92 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/translator_select_test.go @@ -20,6 +20,7 @@ import ( "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/parser" "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/testing/mockdata" "github.com/stretchr/testify/require" + "go.uber.org/zap" "testing" "time" @@ -1086,7 +1087,7 @@ func TestTranslator_TranslateSelectQuerytoBigtable(t *testing.T) { { name: "CqlQuery Without Select Object", query: "UPDATE table_name SET pk1 = 'new_value1', col_int = 'new_value2' WHERE condition;", - wantErr: "mismatched input 'UPDATE' expecting 'SELECT'", + wantErr: "parsing error", want: nil, sessionKeyspace: "test_keyspace", }, @@ -1435,8 +1436,8 @@ func TestTranslator_TranslateSelectQuerytoBigtable(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - tr := NewSelectTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeSelect), tt.sessionKeyspace) + tr := NewSelectTranslator(mockdata.GetSchemaMappingConfig(), zap.NewNop()) + got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeSelect), tt.sessionKeyspace) if tt.wantErr != "" { require.Error(t, err) assert.Contains(t, err.Error(), tt.wantErr) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/types.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/types.go index def9c2a1..325eb67c 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/types.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/select_translator/types.go @@ -3,16 +3,18 @@ package select_translator import ( "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/global/types" schemaMapping "github.com/GoogleCloudPlatform/cloud-bigtable-ecosystem/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/metadata" + "go.uber.org/zap" ) type SelectTranslator struct { schemaMappingConfig *schemaMapping.SchemaMetadata + logger *zap.Logger } func (t *SelectTranslator) QueryType() types.QueryType { return types.QueryTypeSelect } -func NewSelectTranslator(schemaMappingConfig *schemaMapping.SchemaMetadata) types.IQueryTranslator { - return &SelectTranslator{schemaMappingConfig: schemaMappingConfig} +func NewSelectTranslator(schemaMappingConfig *schemaMapping.SchemaMetadata, logger *zap.Logger) types.IQueryTranslator { + return &SelectTranslator{schemaMappingConfig: schemaMappingConfig, logger: logger} } diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/truncate_translator/translator_truncate_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/truncate_translator/translator_truncate_test.go index 43b8f8c7..f5b00d84 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/truncate_translator/translator_truncate_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/truncate_translator/translator_truncate_test.go @@ -94,7 +94,7 @@ func TestTranslateTruncateTableToBigtable(t *testing.T) { t.Run(tt.name, func(t *testing.T) { tr := NewTruncateTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.defaultKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeDelete), tt.defaultKeyspace) if tt.wantErr != "" { require.Error(t, err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/update_translator/translator_update_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/update_translator/translator_update_test.go index ba2f33a6..ba804bf3 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/update_translator/translator_update_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/update_translator/translator_update_test.go @@ -562,7 +562,7 @@ func TestTranslator_TranslateUpdateQuerytoBigtable(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tr := NewUpdateTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeUpdate), tt.sessionKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeUpdate), tt.sessionKeyspace) if tt.wantErr != "" { require.Error(t, err) diff --git a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/use_translator/use_translator_test.go b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/use_translator/use_translator_test.go index fa2d4e71..592defad 100644 --- a/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/use_translator/use_translator_test.go +++ b/cassandra-bigtable-migration-tools/cassandra-bigtable-proxy/translators/use_translator/use_translator_test.go @@ -42,7 +42,7 @@ func TestTranslator_TranslateUpdateQuerytoBigtable(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tr := NewUseTranslator(mockdata.GetSchemaMappingConfig()) - got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.NewParser(tt.query), types.QueryTypeUse), tt.sessionKeyspace) + got, err := tr.Translate(types.NewRawQuery(nil, tt.sessionKeyspace, tt.query, parser.GetParser(tt.query), types.QueryTypeUse), tt.sessionKeyspace) if tt.wantErr != "" { require.Error(t, err)