register executors for bq.source

bruin-data · Jan 10, 2025 · 8c84f0e · 8c84f0e
1 parent 91ca82c
commit 8c84f0e
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 deletions.
diff --git a/cmd/run.go b/cmd/run.go
@@ -638,6 +638,10 @@ func setupExecutors(
 		mainExecutors[pipeline.AssetTypeBigqueryQuery][scheduler.TaskInstanceTypeCustomCheck] = customCheckRunner
 		mainExecutors[pipeline.AssetTypeBigqueryQuery][scheduler.TaskInstanceTypeMetadataPush] = metadataPushOperator
 
+		mainExecutors[pipeline.AssetTypeBigquerySource][scheduler.TaskInstanceTypeMetadataPush] = metadataPushOperator
+		mainExecutors[pipeline.AssetTypeBigquerySource][scheduler.TaskInstanceTypeColumnCheck] = bqCheckRunner
+		mainExecutors[pipeline.AssetTypeBigquerySource][scheduler.TaskInstanceTypeCustomCheck] = customCheckRunner
+
 		mainExecutors[pipeline.AssetTypeBigquerySeed][scheduler.TaskInstanceTypeMain] = seedOperator
 		mainExecutors[pipeline.AssetTypeBigquerySeed][scheduler.TaskInstanceTypeColumnCheck] = bqCheckRunner
 		mainExecutors[pipeline.AssetTypeBigquerySeed][scheduler.TaskInstanceTypeCustomCheck] = customCheckRunner

diff --git a/pkg/bigquery/checks_test.go b/pkg/bigquery/checks_test.go
@@ -74,6 +74,7 @@ func (m *mockQuerierWithResult) IsSameClustering(meta *bigquery.TableMetadata, a
 	args := m.Called(meta, asset)
 	return args.Bool(0)
 }
+
 func (m *mockQuerierWithResult) DeleteTableIfMaterializationTypeMismatch(ctx context.Context, tableName string, asset *pipeline.Asset) error {
 	args := m.Called(ctx, tableName, asset)
 	return args.Error(0)

diff --git a/pkg/executor/defaults.go b/pkg/executor/defaults.go
@@ -30,6 +30,8 @@ var DefaultExecutorsV2 = map[pipeline.AssetType]Config{
 	pipeline.AssetTypeBigquerySource: {
 		scheduler.TaskInstanceTypeMain:         NoOpOperator{},
 		scheduler.TaskInstanceTypeMetadataPush: NoOpOperator{},
+		scheduler.TaskInstanceTypeColumnCheck:  NoOpOperator{},
+		scheduler.TaskInstanceTypeCustomCheck:  NoOpOperator{},
 	},
 	pipeline.AssetTypeBigquerySeed: {
 		scheduler.TaskInstanceTypeMain:         NoOpOperator{},

diff --git a/pkg/snowflake/checks_test.go b/pkg/snowflake/checks_test.go
@@ -44,6 +44,7 @@ func (m *mockQuerierWithResult) Ping(ctx context.Context) error {
 	args := m.Called(ctx)
 	return args.Error(0)
 }
+
 func (m *mockQuerierWithResult) CreateSchemaIfNotExist(ctx context.Context, asset *pipeline.Asset) error {
 	args := m.Called(asset, ctx)
 	return args.Error(0)