diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py index 971be06543ace..1564e57b99771 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py @@ -163,11 +163,6 @@ def __init__( # Set after the job is submitted, parsed from the job submission response. self.job_name: str | None = None - if self.spec_text and (self.spec or self.spec_stage): - raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'") - if not self.spec_text and not (self.spec and self.spec_stage): - raise ValueError("Must provide either 'spec_text' or both 'spec' and 'spec_stage'") - @cached_property def _hook(self) -> SnowflakeHook: return SnowflakeHook( @@ -237,6 +232,11 @@ def on_kill(self) -> None: def execute(self, context: Context) -> str: """Submit and optionally wait for a Snowpark Container Services job.""" + if self.spec_text and (self.spec or self.spec_stage): + raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'") + if not self.spec_text and not (self.spec and self.spec_stage): + raise ValueError("Must provide either 'spec_text' or both 'spec' and 'spec_stage'") + self.job_name = self._submit_job() if not self.job_name: raise RuntimeError("Job name was not returned") diff --git a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py index 3b10bfea06d85..af6114f6616f1 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py @@ -72,8 +72,9 @@ class TestSnowparkContainerJobOperator: ), ) def test_invalid_spec_combinations(self, kwargs, match): + op = _make_operator(**kwargs) with pytest.raises(ValueError, match=match): - _make_operator(**kwargs) + op.execute(context=None) def test_build_sql_with_spec_stage(self): op = _make_operator() diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index 01fd5bc56dbb7..774cb619eeb30 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -80,7 +80,6 @@ providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py: providers/neo4j/src/airflow/providers/neo4j/operators/neo4j.py::Neo4jOperator providers/oracle/src/airflow/providers/oracle/transfers/oracle_to_oracle.py::OracleToOracleOperator providers/papermill/src/airflow/providers/papermill/operators/papermill.py::PapermillOperator -providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py::SnowparkContainerJobOperator providers/ssh/src/airflow/providers/ssh/operators/ssh.py::SSHOperator providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py::SSHRemoteJobOperator providers/standard/src/airflow/providers/standard/operators/bash.py::BashOperator