Merge pull request #420 from bruin-data/patch/fix-metadata-push-for-b…

…q-sensors fix metadata push for bq sensors
bruin-data · Jan 25, 2025 · 023b98e · 023b98e
2 parents 0ecd685 + 6424d3a
commit 023b98e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/cmd/run.go b/cmd/run.go
@@ -638,6 +638,14 @@ func setupExecutors(
 		mainExecutors[pipeline.AssetTypeBigquerySource][scheduler.TaskInstanceTypeColumnCheck] = bqCheckRunner
 		mainExecutors[pipeline.AssetTypeBigquerySource][scheduler.TaskInstanceTypeCustomCheck] = customCheckRunner
 
+		mainExecutors[pipeline.AssetTypeBigqueryTableSensor][scheduler.TaskInstanceTypeMetadataPush] = metadataPushOperator
+		mainExecutors[pipeline.AssetTypeBigqueryTableSensor][scheduler.TaskInstanceTypeColumnCheck] = bqCheckRunner
+		mainExecutors[pipeline.AssetTypeBigqueryTableSensor][scheduler.TaskInstanceTypeCustomCheck] = customCheckRunner
+
+		mainExecutors[pipeline.AssetTypeBigqueryQuerySensor][scheduler.TaskInstanceTypeMetadataPush] = metadataPushOperator
+		mainExecutors[pipeline.AssetTypeBigqueryQuerySensor][scheduler.TaskInstanceTypeColumnCheck] = bqCheckRunner
+		mainExecutors[pipeline.AssetTypeBigqueryQuerySensor][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/executor/defaults.go b/pkg/executor/defaults.go
@@ -9,23 +9,22 @@ type Config map[scheduler.TaskInstanceType]Operator
 
 var DefaultExecutorsV2 = map[pipeline.AssetType]Config{
 	pipeline.AssetTypeBigqueryQuery: {
-		scheduler.TaskInstanceTypeMain:        NoOpOperator{},
-		scheduler.TaskInstanceTypeColumnCheck: NoOpOperator{},
+		scheduler.TaskInstanceTypeMain:         NoOpOperator{},
+		scheduler.TaskInstanceTypeMetadataPush: NoOpOperator{},
+		scheduler.TaskInstanceTypeColumnCheck:  NoOpOperator{},
+		scheduler.TaskInstanceTypeCustomCheck:  NoOpOperator{},
 	},
 	pipeline.AssetTypeBigqueryTableSensor: {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
+		scheduler.TaskInstanceTypeMain:         NoOpOperator{},
+		scheduler.TaskInstanceTypeMetadataPush: NoOpOperator{},
+		scheduler.TaskInstanceTypeColumnCheck:  NoOpOperator{},
+		scheduler.TaskInstanceTypeCustomCheck:  NoOpOperator{},
 	},
 	pipeline.AssetTypeBigqueryQuerySensor: {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
-	},
-	"bq.cost_tracker": {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
-	},
-	"bq.transfer": {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
-	},
-	"bq.sensor.partition": {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
+		scheduler.TaskInstanceTypeMain:         NoOpOperator{},
+		scheduler.TaskInstanceTypeMetadataPush: NoOpOperator{},
+		scheduler.TaskInstanceTypeColumnCheck:  NoOpOperator{},
+		scheduler.TaskInstanceTypeCustomCheck:  NoOpOperator{},
 	},
 	pipeline.AssetTypeBigquerySource: {
 		scheduler.TaskInstanceTypeMain:         NoOpOperator{},
@@ -37,12 +36,6 @@ var DefaultExecutorsV2 = map[pipeline.AssetType]Config{
 		scheduler.TaskInstanceTypeMain:         NoOpOperator{},
 		scheduler.TaskInstanceTypeMetadataPush: NoOpOperator{},
 	},
-	"gcs.from.s3": {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
-	},
-	"gcs.delete": {
-		scheduler.TaskInstanceTypeMain: NoOpOperator{},
-	},
 	"gcs.sensor.object_sensor_with_prefix": {
 		scheduler.TaskInstanceTypeMain: NoOpOperator{},
 	},