Merge pull request #32 from Recidiviz/dan/view-schema-perf

Detect view schema from changed catalog

Author: ohaibbq

internal/contentdata/repository.go

@@ -104,27 +104,50 @@ func (r *Repository) CreateTable(ctx context.Context, tx *connection.Tx, table *
 	return nil
 }
 
-func (r *Repository) CreateView(ctx context.Context, tx *connection.Tx, table *bigqueryv2.Table) error {
+func getSchemaFromResult(result sql.Result) (*bigqueryv2.TableSchema, error) {
+	changedCatalog, err := zetasqlite.ChangedCatalogFromResult(result)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get changed catalog: %w", err)
+	}
+	if len(changedCatalog.Table.Added) != 1 {
+		return nil, fmt.Errorf("catalog detected %d tables added; but expected one", len(changedCatalog.Table.Added))
+	}
+	createdTable := changedCatalog.Table.Added[0]
+	fields := make([]*bigqueryv2.TableFieldSchema, 0, len(createdTable.Columns))
+	for _, col := range createdTable.Columns {
+		zetasqlType, err := col.Type.ToZetaSQLType()
+		if err != nil {
+			return nil, fmt.Errorf("failed to get zetasql type: %w", err)
+		}
+		fields = append(fields, types.TableFieldSchemaFromZetaSQLType(col.Name, zetasqlType))
+	}
+
+	return &bigqueryv2.TableSchema{Fields: fields}, nil
+}
+
+func (r *Repository) CreateView(ctx context.Context, tx *connection.Tx, table *bigqueryv2.Table) (*bigqueryv2.TableSchema, error) {
 	if err := tx.ContentRepoMode(); err != nil {
-		return err
+		return nil, err
 	}
 	defer func() {
 		_ = tx.MetadataRepoMode()
 	}()
 	ref := table.TableReference
 	if ref == nil {
-		return fmt.Errorf("TableReference is nil")
+		return nil, fmt.Errorf("TableReference is nil")
 	}
 	viewDefinition := table.View
 	if viewDefinition == nil {
-		return fmt.Errorf("ViewDefinition is nil")
+		return nil, fmt.Errorf("ViewDefinition is nil")
 	}
 	tablePath := r.tablePath(ref.ProjectId, ref.DatasetId, ref.TableId)
 	query := fmt.Sprintf("CREATE VIEW `%s` AS (%s)", tablePath, strings.TrimRight(viewDefinition.Query, ViewQueryEndCutset))
-	if _, err := tx.Tx().ExecContext(ctx, query); err != nil {
-		return fmt.Errorf("failed to create view %s: %w", query, err)
+	result, err := tx.Tx().ExecContext(ctx, query)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create view %s: %w", query, err)
 	}
-	return nil
+	schema, err := getSchemaFromResult(result)
+	return schema, err
 }
 
 func (r *Repository) encodeSchemaField(field *bigqueryv2.TableFieldSchema) string {

server/handler.go

@@ -2644,18 +2644,6 @@ func createTableMetadata(ctx context.Context, tx *connection.Tx, server *Server,
 	if table.View != nil {
 		table.Type = string(internaltypes.ViewTableType)
 		table.View.Query = strings.TrimRight(table.View.Query, contentdata.ViewQueryEndCutset)
-		response, err := server.contentRepo.Query(
-			ctx,
-			tx,
-			table.TableReference.ProjectId,
-			table.TableReference.DatasetId,
-			fmt.Sprintf("SELECT * FROM (%s) LIMIT 0", table.View.Query),
-			nil,
-		)
-		if err != nil {
-			return nil, errInvalidQuery(fmt.Sprintf("could not determine schema from query: %s", err.Error()))
-		}
-		table.Schema = response.Schema
 	}
 	table.Kind = "bigquery#table"
 	table.SelfLink = fmt.Sprintf(
@@ -2703,20 +2691,27 @@ func (h *tablesInsertHandler) Handle(ctx context.Context, r *tablesInsertRequest
 	}
 	defer tx.RollbackIfNotCommitted()
 
-	table, serverErr := createTableMetadata(ctx, tx, r.server, r.project, r.dataset, r.table)
-	if serverErr != nil {
-		return nil, serverErr
-	}
-	if table.Type == string(internaltypes.DefaultTableType) && table.Schema != nil {
-		if err := r.server.contentRepo.CreateTable(ctx, tx, r.table); err != nil {
+	if r.table.View != nil {
+		schema, err := r.server.contentRepo.CreateView(ctx, tx, r.table)
+		if err != nil {
+			if strings.HasSuffix(err.Error(), "already exists (1)") {
+				return nil, errDuplicate(err.Error())
+			}
 			return nil, errInvalidQuery(err.Error())
 		}
-	}
-	if table.Type == string(internaltypes.ViewTableType) {
-		if err := r.server.contentRepo.CreateView(ctx, tx, r.table); err != nil {
+		r.table.Schema = schema
+	} else if r.table.Schema != nil {
+		if err := r.server.contentRepo.CreateTable(ctx, tx, r.table); err != nil {
+			if strings.HasSuffix(err.Error(), "already exists (1)") {
+				return nil, errDuplicate(err.Error())
+			}
 			return nil, errInvalidQuery(err.Error())
 		}
 	}
+	table, serverErr := createTableMetadata(ctx, tx, r.server, r.project, r.dataset, r.table)
+	if serverErr != nil {
+		return nil, serverErr
+	}
 	if err := tx.Commit(); err != nil {
 		return nil, errInternalError(fmt.Errorf("failed to commit table: %w", err).Error())
 	}

server/server_test.go

@@ -1223,7 +1223,7 @@ func TestDuplicateTableWithSchema(t *testing.T) {
 			t.Fatalf("%+v", ge)
 		}
 	} else {
-		t.Fatalf(("Threre should be error, when table name duplicates."))
+		t.Fatalf(("There should be error, when table name duplicates."))
 	}
 }