From 1c5ff4ea8fd24d78e0aff7951a85fe6d69e751cf Mon Sep 17 00:00:00 2001 From: Nathan Villaescusa Date: Mon, 22 Apr 2024 22:07:01 -0700 Subject: [PATCH 1/2] [bigquery] Clean up `Store.putTable` --- clients/bigquery/bigquery.go | 44 ++++--------------------------- clients/bigquery/bigquery_test.go | 21 --------------- 2 files changed, 5 insertions(+), 60 deletions(-) diff --git a/clients/bigquery/bigquery.go b/clients/bigquery/bigquery.go index 5de18e30f..b79ac7f47 100644 --- a/clients/bigquery/bigquery.go +++ b/clients/bigquery/bigquery.go @@ -5,7 +5,6 @@ import ( "fmt" "log/slog" "os" - "strings" "cloud.google.com/go/bigquery" _ "github.com/viant/bigquery" @@ -76,7 +75,7 @@ func (s *Store) PrepareTemporaryTable(tableData *optimization.TableData, tableCo } // Load the data - return s.putTable(context.Background(), tableData.TopicConfig().Database, tempTableID, rows) + return s.putTable(context.Background(), tempTableID, rows) } func (s *Store) IdentifierFor(topicConfig kafkalib.TopicConfig, table string) types.TableIdentifier { @@ -124,52 +123,19 @@ func (s *Store) GetClient(ctx context.Context) *bigquery.Client { return client } -func tableRelName(fqName string) (string, error) { - fqNameParts := strings.Split(fqName, ".") - if len(fqNameParts) < 3 { - return "", fmt.Errorf("invalid fully qualified name: %s", fqName) - } - - return strings.Join(fqNameParts[2:], "."), nil -} - -func (s *Store) putTable(ctx context.Context, dataset string, tableID types.TableIdentifier, rows []*Row) error { +func (s *Store) putTable(ctx context.Context, tableID types.TableIdentifier, rows []*Row) error { bqTableID, ok := tableID.(TableIdentifier) if !ok { - return fmt.Errorf("unable to cast tableID to BigQuery TableIdentifier") - } - - tableName := tableID.FullyQualifiedName() - relTableName, err := tableRelName(tableName) - if err != nil { - return fmt.Errorf("failed to get table name: %w", err) - } - - if dataset != bqTableID.Dataset() { - // TODO: [tableID] has [Dataset] on it, don't need to pass it along. - slog.Error("BigQuery dataset is different", - slog.String("dataset", dataset), - slog.String("bqTableID.Dataset", bqTableID.Dataset()), - slog.String("fqn", tableName), - ) - } - - if relTableName != bqTableID.Table() { - // TODO: Use [bqTableID.Table] instead of [relTableName]. - slog.Error("BigQuery table name is different", - slog.String("relTableName", relTableName), - slog.String("bqTableID.Table", bqTableID.Table()), - slog.String("fqn", tableName), - ) + return fmt.Errorf("uanble to cast tableID to a BigQuery TableIdentifier") } client := s.GetClient(ctx) defer client.Close() batch := NewBatch(rows, s.batchSize) - inserter := client.Dataset(dataset).Table(relTableName).Inserter() + inserter := client.Dataset(bqTableID.Dataset()).Table(bqTableID.Table()).Inserter() for batch.HasNext() { - if err = inserter.Put(ctx, batch.NextChunk()); err != nil { + if err := inserter.Put(ctx, batch.NextChunk()); err != nil { return fmt.Errorf("failed to insert rows: %w", err) } } diff --git a/clients/bigquery/bigquery_test.go b/clients/bigquery/bigquery_test.go index c6b6d25c9..2a6fca484 100644 --- a/clients/bigquery/bigquery_test.go +++ b/clients/bigquery/bigquery_test.go @@ -13,27 +13,6 @@ import ( "github.com/stretchr/testify/assert" ) -func (b *BigQueryTestSuite) TestTableRelName() { - { - relName, err := tableRelName("project.dataset.table") - assert.NoError(b.T(), err) - assert.Equal(b.T(), "table", relName) - } - { - relName, err := tableRelName("project.dataset.table.table") - assert.NoError(b.T(), err) - assert.Equal(b.T(), "table.table", relName) - } - { - // All the possible errors - _, err := tableRelName("project.dataset") - assert.ErrorContains(b.T(), err, "invalid fully qualified name: project.dataset") - - _, err = tableRelName("project") - assert.ErrorContains(b.T(), err, "invalid fully qualified name: project") - } -} - func TestTempTableName(t *testing.T) { trimTTL := func(tableName string) string { lastUnderscore := strings.LastIndex(tableName, "_") From c80ce2208fc64629605d71fcb11e20931f1b3c52 Mon Sep 17 00:00:00 2001 From: Nathan Villaescusa Date: Mon, 22 Apr 2024 22:08:36 -0700 Subject: [PATCH 2/2] Revert --- clients/bigquery/bigquery.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/clients/bigquery/bigquery.go b/clients/bigquery/bigquery.go index b79ac7f47..f2ba0182d 100644 --- a/clients/bigquery/bigquery.go +++ b/clients/bigquery/bigquery.go @@ -126,7 +126,7 @@ func (s *Store) GetClient(ctx context.Context) *bigquery.Client { func (s *Store) putTable(ctx context.Context, tableID types.TableIdentifier, rows []*Row) error { bqTableID, ok := tableID.(TableIdentifier) if !ok { - return fmt.Errorf("uanble to cast tableID to a BigQuery TableIdentifier") + return fmt.Errorf("unable to cast tableID to BigQuery TableIdentifier") } client := s.GetClient(ctx)