Adjust test for Airflow-2.9 (#2149)

- pin google-cloud-bigquery
- Remove conditional query building for different Airflow version
- Adjust dataset and cleanoperator test 
- clean checkoperator test

python-sdk/pyproject.toml

@@ -57,7 +57,8 @@ google = [
     "protobuf",
     "apache-airflow-providers-google>=10.15.0",
     "sqlalchemy-bigquery>=1.3.0",
-    "smart-open[gcs]>=5.2.1,<7.0.0"
+    "smart-open[gcs]>=5.2.1,<7.0.0",
+    "google-cloud-bigquery<3.21.0"
 ]
 snowflake = [
     "apache-airflow-providers-snowflake>=5.3.0",
@@ -126,7 +127,8 @@ all = [
     "azure-storage-blob",
     "apache-airflow-providers-microsoft-mssql>=3.2",
     "airflow-provider-duckdb>=0.0.2",
-    "apache-airflow-providers-mysql"
+    "apache-airflow-providers-mysql",
+    "google-cloud-bigquery<3.21.0"
 ]
 doc = [
     "myst-parser>=0.17",

python-sdk/src/astro/sql/operators/data_validations/check_table.py

@@ -56,12 +56,7 @@ def execute(self, context: "Context"):
         db = create_database(self.dataset.conn_id)
         self.table = db.get_table_qualified_name(self.dataset)
         self.conn_id = self.dataset.conn_id
-        # apache-airflow-providers-common-sql == 1.2.0 which is compatible with airflow 2.2.5 implements the self.sql
-        # differently compared to apache-airflow-providers-common-sql == 1.3.3
-        try:
-            self.sql = f"SELECT check_name, check_result FROM ({self._generate_sql_query()}) AS check_table"
-        except AttributeError:
-            self.sql = f"SELECT * FROM {self.table};"
+        self.sql = f"SELECT check_name, check_result FROM ({self._generate_sql_query()}) AS check_table"
         super().execute(context)
 
     def get_db_hook(self) -> Any:

python-sdk/tests/airflow_tests/test_datasets.py

@@ -104,7 +104,6 @@ def test_kwargs_with_temp_table():
 @pytest.mark.skipif(airflow.__version__ < "2.4.0", reason="Require Airflow version >= 2.4.0")
 def test_example_dataset_dag():
     from airflow.datasets import Dataset
-    from airflow.models.dataset import DatasetModel
 
     dir_path = os.path.dirname(os.path.realpath(__file__))
     db = DagBag(dir_path + "/../../example_dags/example_datasets.py")
@@ -115,9 +114,8 @@ def test_example_dataset_dag():
     outlets = producer_dag.tasks[-1].outlets
     assert isinstance(outlets[0], Dataset)
     # Test that dataset_triggers is only set if all the instances passed to the DAG object are Datasets
-    assert consumer_dag.dataset_triggers == outlets
+    assert consumer_dag.dataset_triggers.objects[0] == outlets[0]
     assert outlets[0].uri == "astro://postgres_conn@?table=imdb_movies"
-    assert DatasetModel.from_public(outlets[0]) == Dataset("astro://postgres_conn@?table=imdb_movies")
 
 
 def test_disable_auto_inlets_outlets():

python-sdk/tests/sql/operators/test_cleanup.py

@@ -105,19 +105,20 @@ def test_error_raised_with_blocking_op_executors(
     reason="BackfillJobRunner and Job classes are only available in airflow >= 2.6",
 )
 @pytest.mark.parametrize(
-    "executor_in_job,executor_in_cfg,expected_val",
+    "executor_in_job, executor_in_cfg, expected_val",
     [
-        (SequentialExecutor(), "LocalExecutor", True),
+        (SequentialExecutor(), "SequentialExecutor", True),
         (LocalExecutor(), "LocalExecutor", False),
         (None, "LocalExecutor", False),
         (None, "SequentialExecutor", True),
     ],
 )
-def test_single_worker_mode_backfill(executor_in_job, executor_in_cfg, expected_val):
+def test_single_worker_mode_backfill(monkeypatch, executor_in_job, executor_in_cfg, expected_val):
     """Test that if we run Backfill Job it should be marked as single worker node"""
     from airflow.jobs.backfill_job_runner import BackfillJobRunner
     from airflow.jobs.job import Job
 
+    monkeypatch.setattr("airflow.executors.executor_loader._executor_names", [])
     dag = DAG("test_single_worker_mode_backfill", start_date=datetime(2022, 1, 1))
     dr = DagRun(dag_id=dag.dag_id)
 
@@ -175,17 +176,18 @@ def test_single_worker_mode_backfill_airflow_2_5(executor_in_job, executor_in_cf
 @pytest.mark.parametrize(
     "executor_in_job,executor_in_cfg,expected_val",
     [
-        (SequentialExecutor(), "LocalExecutor", True),
+        (SequentialExecutor(), "SequentialExecutor", True),
         (LocalExecutor(), "LocalExecutor", False),
         (None, "LocalExecutor", False),
         (None, "SequentialExecutor", True),
     ],
 )
-def test_single_worker_mode_scheduler_job(executor_in_job, executor_in_cfg, expected_val):
+def test_single_worker_mode_scheduler_job(monkeypatch, executor_in_job, executor_in_cfg, expected_val):
     """Test that if we run Scheduler Job it should be marked as single worker node"""
     from airflow.jobs.job import Job
     from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
 
+    monkeypatch.setattr("airflow.executors.executor_loader._executor_names", [])
     dag = DAG("test_single_worker_mode_scheduler_job", start_date=datetime(2022, 1, 1))
     dr = DagRun(dag_id=dag.dag_id)
 

python-sdk/tests_integration/sql/operators/data_validation/test_check_column.py

@@ -5,7 +5,6 @@
 from astro import sql as aql
 from astro.constants import Database
 from astro.files import File
-from astro.table import Table
 from tests.sql.operators import utils as test_utils
 
 CWD = pathlib.Path(__file__).parent
@@ -22,7 +21,6 @@
         {
             "database": Database.BIGQUERY,
             "file": File(path=str(CWD) + "/../../../data/data_validation.csv"),
-            "table": Table(conn_id="gcp_conn_project"),
         },
         {
             "database": Database.POSTGRES,

python-sdk/tests_integration/sql/operators/data_validation/test_check_table.py

@@ -5,7 +5,6 @@
 from astro import sql as aql
 from astro.constants import Database
 from astro.files import File
-from astro.table import Table
 from tests.sql.operators import utils as test_utils
 
 CWD = pathlib.Path(__file__).parent
@@ -22,7 +21,6 @@
         {
             "database": Database.BIGQUERY,
             "file": File(path=str(CWD) + "/../../../data/homes_main.csv"),
-            "table": Table(conn_id="gcp_conn_project"),
         },
         {
             "database": Database.POSTGRES,