Snowflake: support warehouse switching (#2492)

Author: m1n0

soda-core/src/soda_core/common/data_source_impl.py

@@ -235,3 +235,11 @@ def verify_if_table_exists(self, prefixes: list[str], table_name: str) -> bool:
             fully_qualified_table_name.table_name == table_name
             for fully_qualified_table_name in fully_qualified_table_names
         )
+
+    def switch_warehouse(self, warehouse: str) -> None:
+        # Noop by default, only some data sources need to implement this
+        pass
+
+    def get_current_warehouse(self) -> Optional[str]:
+        # Noop by default, only some data sources need to implement this
+        return None

soda-core/src/soda_core/common/soda_cloud_dto.py

@@ -46,6 +46,12 @@ class DatasetConfigurationsDTO(BaseModel):
     dataset_configurations: list[DatasetConfigurationDTO] = Field(..., alias="datasetConfigurations")
 
 
+class ComputeWarehouseOverrideDTO(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    name: str = Field(..., alias="name")
+
+
 class DatasetConfigurationDTO(BaseModel):
     model_config = ConfigDict(populate_by_name=True, extra="allow")
 
@@ -61,6 +67,7 @@ class DatasetConfigurationDTO(BaseModel):
     test_row_sampler_configuration: Optional[TestRowSamplerConfigurationDTO] = Field(
         None, alias="testRowSamplerConfiguration"
     )
+    compute_warehouse_override: Optional[ComputeWarehouseOverrideDTO] = Field(None, alias="computeWarehouseOverride")
 
 
 class TestRowSamplerAbsoluteLimitDTO(BaseModel):

soda-core/src/soda_core/contracts/impl/contract_verification_impl.py

@@ -410,11 +410,19 @@ def __init__(
         # TODO replace usage of self.sql_qualified_dataset_name with self.dataset_identifier
         self.sql_qualified_dataset_name: Optional[str] = None
 
+        self.datasource_warehouse: Optional[str] = None
+        self.compute_warehouse: Optional[str] = None
+
         if data_source_impl:
             # TODO replace usage of self.sql_qualified_dataset_name with self.dataset_identifier
             self.sql_qualified_dataset_name = data_source_impl.sql_dialect.qualify_dataset_name(
                 dataset_prefix=self.dataset_prefix, dataset_name=self.dataset_name
             )
+            if hasattr(data_source_impl.data_source_model, "warehouse"):
+                self.datasource_warehouse = data_source_impl.data_source_model.warehouse
+
+            if self.datasource_warehouse is None:
+                self.datasource_warehouse = data_source_impl.get_current_warehouse()
 
         from soda_core.contracts.impl.check_types.row_count_check import (
             RowCountMetricImpl,
@@ -450,10 +458,13 @@ def __init__(
                 self.sampler_type = self.dataset_configuration.test_row_sampler_configuration.test_row_sampler.type
                 self.sampler_limit = self.dataset_configuration.test_row_sampler_configuration.test_row_sampler.limit
 
+            if self.dataset_configuration.compute_warehouse_override:
+                self.compute_warehouse = self.dataset_configuration.compute_warehouse_override.name
+
         if self.should_apply_sampling:
             logger.info(
                 f"Row sampling is enabled for dataset {self.dataset_identifier.to_string()} "
-                f"with sampler {self.dataset_configuration.test_row_sampler_configuration.test_row_sampler.type}"
+                f"with sampler config: type:'{self.dataset_configuration.test_row_sampler_configuration.test_row_sampler.type}', limit:'{self.dataset_configuration.test_row_sampler_configuration.test_row_sampler.limit}'"
             )
 
             # This modifies the CTE to include sampling by accessing the first element of the cte_query list, may be flaky. Consider adding a better way to modify queries, or change AST to a 3rd party library which may support it already.
@@ -598,6 +609,16 @@ def _parse_columns(self, contract_yaml: ContractYaml) -> list[ColumnImpl]:
         return columns
 
     def verify(self) -> ContractVerificationResult:
+        if (
+            self.data_source_impl
+            and self.datasource_warehouse
+            and self.compute_warehouse
+            and self.datasource_warehouse != self.compute_warehouse
+        ):
+            logger.info(
+                f"Switching warehouse from '{self.datasource_warehouse}' to '{self.compute_warehouse}' for Contract verification of dataset '{self.dataset_identifier.to_string()}'"
+            )
+            self.data_source_impl.switch_warehouse(self.compute_warehouse)
         data_source: Optional[DataSource] = None
         check_results: list[CheckResult] = []
         measurements: list[Measurement] = []
@@ -724,6 +745,18 @@ def verify(self) -> ContractVerificationResult:
                     scan_id=scan_id, exc=e, contract_verification_handler=contract_verification_handler
                 )
 
+        # Switch back to original warehouse if changed
+        if (
+            self.data_source_impl
+            and self.compute_warehouse
+            and self.datasource_warehouse
+            and self.datasource_warehouse != self.compute_warehouse
+        ):
+            logger.info(
+                f"Switching back warehouse to '{self.datasource_warehouse}' after Contract verification of dataset '{self.dataset_identifier.to_string()}'"
+            )
+            self.data_source_impl.switch_warehouse(self.datasource_warehouse)
+
         return contract_verification_result
 
     def __get_dataset_id(self, soda_cloud_response_json: dict, qualified_dataset_name: str) -> Optional[str]:

soda-snowflake/src/soda_snowflake/common/data_sources/snowflake_data_source.py

@@ -36,6 +36,19 @@ def _create_data_source_connection(self) -> DataSourceConnection:
             name=self.data_source_model.name, connection_properties=self.data_source_model.connection_properties
         )
 
+    def switch_warehouse(self, warehouse: str) -> None:
+        switch_warehouse_sql = f"USE WAREHOUSE {warehouse}"
+        self.execute_query(switch_warehouse_sql)
+
+    def get_current_warehouse(self) -> Optional[str]:
+        sql = "SELECT CURRENT_WAREHOUSE()"
+        current_warehouse_sql = "SELECT CURRENT_WAREHOUSE()"
+        result = self.execute_query(current_warehouse_sql)
+        result_rows = result.rows
+        row = result_rows[0] if result_rows else None
+
+        return row[0] if row and row[0] else None
+
 
 class SnowflakeSqlDialect(SqlDialect):
     SODA_DATA_TYPE_SYNONYMS = (