Enable unit tests for Python-3.11

.github/workflows/ci-python-sdk.yaml

@@ -178,6 +178,8 @@ jobs:
             ~/.cache/pip
             .nox
           key: ${{ runner.os }}-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
+      - run: pip --version
       - run: cat ../.github/ci-test-connections.yaml > test-connections.yaml
       - run: python -c 'import os; print(os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", "").strip())' > ${{ env.GOOGLE_APPLICATION_CREDENTIALS }}
       - run: sqlite3 /tmp/sqlite_default.db "VACUUM;"
@@ -203,10 +205,10 @@ jobs:
       AZURE_WASB_CONN_STRING: ${{ secrets.AZURE_WASB_CONN_STRING }}
       AZURE_WASB_ACCESS_KEY: ${{ secrets.AZURE_WASB_ACCESS_KEY }}
 
-  Run-Unit-tests-Airflow-2-5:
+  Run-Unit-tests-Airflow-2-6:
     strategy:
       matrix:
-        version: [ '3.8', '3.9', '3.10' ]
+        version: [ '3.8', '3.9', '3.10', '3.11' ]
     if: >-
       github.event_name == 'push' ||
       (
@@ -240,9 +242,13 @@ jobs:
             ~/.cache/pip
             .nox
           key: ${{ runner.os }}-2.5-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
+      - run: sudo add-apt-repository "deb http://dl.bintray.com/apache/arrow/ubuntu $(lsb_release --codename --short) main" |
+             sudo apt-get install libarrow-dev libparquet-dev libarrow-python-dev
+      - run: pip3 --version
       - run: sqlite3 /tmp/sqlite_default.db "VACUUM;"
       - run: pip3 install nox
-      - run: nox -s "test-${{ matrix.version }}(airflow='2.6.0')" -- tests/ --cov=src --cov-report=xml --cov-branch
+      - run: nox -s "test-${{ matrix.version }}(airflow='2.6.3')" -- tests/ --cov=src --cov-report=xml --cov-branch
       - name: Upload coverage
         uses: actions/upload-artifact@v2
         with:
@@ -315,11 +321,12 @@ jobs:
             ~/.cache/pip
             .nox
           key: ${{ runner.os }}-2.5-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
       - run: cat ../.github/ci-test-connections.yaml > test-connections.yaml
       - run: python -c 'import os; print(os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", "").strip())' > ${{ env.GOOGLE_APPLICATION_CREDENTIALS }}
       - run: sqlite3 /tmp/sqlite_default.db "VACUUM;"
       - run: pip3 install nox
-      - run: nox -s "test-3.8(airflow='2.6.0')" -- tests_integration/ -k "test_load_file.py and not redshift" --splits 3 --group ${{ matrix.group }} --store-durations --durations-path /tmp/durations-${{ matrix.group }} --cov=src --cov-report=xml --cov-branch
+      - run: nox -s "test-3.8(airflow='2.6.3')" -- tests_integration/ -k "test_load_file.py and not redshift" --splits 3 --group ${{ matrix.group }} --store-durations --durations-path /tmp/durations-${{ matrix.group }} --cov=src --cov-report=xml --cov-branch
       - run: cat /tmp/durations-${{ matrix.group }}
       - name: Upload coverage
         uses: actions/upload-artifact@v2
@@ -411,11 +418,12 @@ jobs:
             ~/.cache/pip
             .nox
           key: ${{ runner.os }}-2.5-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
       - run: cat ../.github/ci-test-connections.yaml > test-connections.yaml
       - run: python -c 'import os; print(os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", "").strip())' > ${{ env.GOOGLE_APPLICATION_CREDENTIALS }}
       - run: sqlite3 /tmp/sqlite_default.db "VACUUM;"
       - run: pip3 install nox
-      - run: nox -s "test-3.8(airflow='2.6.0')" -- tests_integration/ -k "test_example_dags.py and not redshift" --splits 3 --group ${{ matrix.group }} --store-durations --durations-path /tmp/durations-${{ matrix.group }} --cov=src --cov-report=xml --cov-branch
+      - run: nox -s "test-3.8(airflow='2.6.3')" -- tests_integration/ -k "test_example_dags.py and not redshift" --splits 3 --group ${{ matrix.group }} --store-durations --durations-path /tmp/durations-${{ matrix.group }} --cov=src --cov-report=xml --cov-branch
       - run: cat /tmp/durations-${{ matrix.group }}
       - name: Upload coverage
         uses: actions/upload-artifact@v2
@@ -507,11 +515,12 @@ jobs:
             ~/.cache/pip
             .nox
           key: ${{ runner.os }}-2.5-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
       - run: cat ../.github/ci-test-connections.yaml > test-connections.yaml
       - run: python -c 'import os; print(os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", "").strip())' > ${{ env.GOOGLE_APPLICATION_CREDENTIALS }}
       - run: sqlite3 /tmp/sqlite_default.db "VACUUM;"
       - run: pip3 install nox
-      - run: nox -s "test-3.8(airflow='2.6.0')" -- tests_integration/ -k "not test_load_file.py and not test_example_dags.py and not redshift" --splits 11 --group ${{ matrix.group }} --store-durations --durations-path /tmp/durations-${{ matrix.group }} --cov=src --cov-report=xml --cov-branch
+      - run: nox -s "test-3.8(airflow='2.6.3')" -- tests_integration/ -k "not test_load_file.py and not test_example_dags.py and not redshift" --splits 11 --group ${{ matrix.group }} --store-durations --durations-path /tmp/durations-${{ matrix.group }} --cov=src --cov-report=xml --cov-branch
       - run: cat /tmp/durations-${{ matrix.group }}
       - name: Upload coverage
         uses: actions/upload-artifact@v2
@@ -601,6 +610,7 @@ jobs:
             ~/.cache/pip
             .nox
           key: ${{ runner.os }}-2.2.5-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
       - run: cat ../.github/ci-test-connections.yaml > test-connections.yaml
       - run: python -c 'import os; print(os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", "").strip())' > ${{ env.GOOGLE_APPLICATION_CREDENTIALS }}
       - run: sqlite3 /tmp/sqlite_default.db "VACUUM;"
@@ -627,7 +637,7 @@ jobs:
       fail-fast: false
       matrix:
         python: [ '3.7', '3.8', '3.9', '3.10' ]
-        airflow: [ '2.2.5', '2.3.4', '2.4.2', '2.5.3', '2.6.0']
+        airflow: [ '2.2.5', '2.3.4', '2.4.2', '2.5.3', '2.6.3']
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
@@ -643,6 +653,7 @@ jobs:
         with:
           path: ${{ steps.pip-cache.outputs.pip_cache_dir }}
           key: constraints-${{ matrix.python }}-${{ matrix.airflow }}-${{ hashFiles('python-sdk/pyproject.toml') }}-${{ hashFiles('python-sdk/src/astro/__init__.py') }}
+      - run: sudo apt-get install -y --no-install-recommends build-essential libsasl2-2 libsasl2-dev libsasl2-modules freetds-dev libssl-dev libkrb5-dev libmariadb-dev
       - run: pip3 install -U pip wheel build nox
       - run: nox -s build
       - run: pip3 install 'apache-airflow==${{ matrix.airflow }}' "$(ls dist/*.whl)[all]"
@@ -658,7 +669,7 @@ jobs:
   Code-Coverage:
     if: github.event.action != 'labeled'
     needs:
-      - Run-Unit-tests-Airflow-2-5
+      - Run-Unit-tests-Airflow-2-6
       - Run-Integration-tests-Airflow-2-5
       - Run-load-file-Integration-Airflow-2-5
       - Run-example-dag-Integration-Airflow-2-5
@@ -693,7 +704,7 @@ jobs:
   collect-durations:
     runs-on: ubuntu-latest
     needs:
-      - Run-Unit-tests-Airflow-2-5
+      - Run-Unit-tests-Airflow-2-6
       - Run-Integration-tests-Airflow-2-5
       - Run-load-file-Integration-Airflow-2-5
     steps:
@@ -723,7 +734,7 @@ jobs:
     if: github.event_name == 'release'
     name: Build and publish Python 🐍 distributions 📦 to PyPI
     needs:
-      - Run-Unit-tests-Airflow-2-5
+      - Run-Unit-tests-Airflow-2-6
       - Run-example-dag-tests-Airflow-2-2-5
       - Run-Integration-tests-Airflow-2-5
       - Run-load-file-Integration-Airflow-2-5

python-sdk/example_dags/calculate_popular_movies_by_genre.py

@@ -31,6 +31,6 @@ def top_movies_by_genre(input_df: pd.DataFrame):
     )
     movies_by_genre = top_movies_by_genre(input_df=imdb_movies)
     with tempfile.NamedTemporaryFile(suffix=".csv") as tmp_file:
-        ExportToFileOperator.partial(output_file=File(path=tmp_file.name), task_id="export_movies").expand(
-            input_data=movies_by_genre
-        )
+        ExportToFileOperator.partial(
+            if_exists="replace", output_file=File(path=tmp_file.name), task_id="export_movies"
+        ).expand(input_data=movies_by_genre)

python-sdk/noxfile.py

@@ -19,8 +19,8 @@ def dev(session: nox.Session) -> None:
     session.install("-e", ".[all,tests]")
 
 
-@nox.session(python=["3.7", "3.8", "3.9", "3.10"])
-@nox.parametrize("airflow", ["2.2.5", "2.4", "2.5.3", "2.6.0"])
+@nox.session(python=["3.7", "3.8", "3.9", "3.10", "3.11"])
+@nox.parametrize("airflow", ["2.2.5", "2.4", "2.5.3", "2.6.3"])
 def test(session: nox.Session, airflow) -> None:
     """Run both unit and integration tests."""
     env = {
@@ -37,11 +37,29 @@ def test(session: nox.Session, airflow) -> None:
         # update the constraints file used below with that  version of provider
         # For example as part of MSSQL support we need apache-airflow-providers-microsoft-mssql>=3.2 and this
         # has been updated in the below constraint file.
-        session.install(f"apache-airflow=={airflow}", "-c", "tests/modified_constraint_file.txt")
-        session.install("-e", ".[all,tests]", "-c", "tests/modified_constraint_file.txt")
+        session.install(f"apache-airflow=={airflow}", "-c", "tests/modified_constraint_file_2.2.5.txt")
+        session.install("-e", ".[all,tests]", "-c", "tests/modified_constraint_file_2.2.5.txt")
         session.install("apache-airflow-providers-common-sql==1.2.0")
         # install smart-open 6.3.0 since it has FTP implementation
         session.install("smart-open>=6.3.0")
+    elif airflow == "2.6.3":
+        env["AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES"] = "airflow\\.* astro\\.*"
+
+        session.install(
+            f"apache-airflow=={airflow}",
+            "-c",
+            f"tests/modified_constraint_file_2.6.3-python-{session.python}.txt",
+        )
+        if session.python == "3.11":
+            session.install(
+                "-e",
+                ".[all,tests]",
+                "-c",
+                f"tests/modified_constraint_file_2.6.3-python-{session.python}.txt",
+            )
+        else:
+            session.install("-e", ".[all,tests]")
+
     else:
         env["AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES"] = "airflow\\.* astro\\.*"
 
@@ -151,7 +169,7 @@ def build_docs(session: nox.Session) -> None:
 
 
 @nox.session(python=["3.7", "3.8", "3.9", "3.10"])
-@nox.parametrize("airflow", ["2.2.5", "2.3.4", "2.4.2", "2.5.3", "2.6.0"])
+@nox.parametrize("airflow", ["2.2.5", "2.3.4", "2.4.2", "2.5.3", "2.6.3"])
 def generate_constraints(session: nox.Session, airflow) -> None:
     """Generate constraints file"""
     session.install("wheel")

python-sdk/pyproject.toml

@@ -17,15 +17,16 @@ license = { file = "LICENSE" }
 
 requires-python = ">=3.7"
 dependencies = [
-    "apache-airflow>=2.0",
+    "apache-airflow==2.6.3",
     "attrs>=20.0",
     "pandas>=1.3.4,<2.0.0", # Pinning it to <2.0.0 to avoid breaking changes
     "pyarrow",
     "pydantic>=1.10.0,<2.0.0", # Airflow & Pydantic issue: https://github.com/apache/airflow/issues/32311
     "python-frontmatter",
     "smart-open",
     "SQLAlchemy>=1.3.18",
-    "cached_property>=1.5.0;python_version<='3.7'"
+    "cached_property>=1.5.0;python_version<='3.7'",
+    #"pandas_gbq>=0.15.0"
 ]
 
 keywords = ["airflow", "provider", "astronomer", "sql", "decorator", "task flow", "elt", "etl", "dag"]
@@ -55,15 +56,15 @@ tests = [
     "sqlalchemy-stubs", # Change when sqlalchemy is upgraded https://docs.sqlalchemy.org/en/14/orm/extensions/mypy.html
 ]
 google = [
-    "protobuf<=3.20", # Google bigquery client require protobuf <= 3.20.0. We can remove the limitation when this limitation is removed
+    #"protobuf<=3.20", # Google bigquery client require protobuf <= 3.20.0. We can remove the limitation when this limitation is removed
     "apache-airflow-providers-google>=6.4.0",
     "sqlalchemy-bigquery>=1.3.0",
     "smart-open[gcs]>=5.2.1"
 ]
 snowflake = [
     "apache-airflow-providers-snowflake",
     "snowflake-sqlalchemy>=1.2.0",
-    "snowflake-connector-python[pandas]<3.0.0",
+    "snowflake-connector-python[pandas]",
     # pinning snowflake-connector-python[pandas]<3.0.0 due to a conflict in snowflake-connector-python/pyarrow/google
     # packages and pandas-gbq/google packages which is forcing pandas-gbq of version 0.13.2 installed, which is not
     # compatible with pandas 1.5.3
@@ -114,7 +115,7 @@ all = [
     "apache-airflow-providers-snowflake",
     "apache-airflow-providers-sftp",
     "smart-open[all]>=5.2.1",
-    "snowflake-connector-python[pandas]<3.0.0",
+    "snowflake-connector-python[pandas]",
     # pinning snowflake-connector-python[pandas]<3.0.0 due to a conflict in snowflake-connector-python/pyarrow/google
     # packages and pandas-gbq/google packages which is forcing pandas-gbq of version 0.13.2 installed, which is not
     # compatible with pandas 1.5.3
@@ -123,7 +124,7 @@ all = [
     "databricks-cli",
     "apache-airflow-providers-databricks",
     "s3fs",
-    "protobuf<=3.20", # Google bigquery client require protobuf <= 3.20.0. We can remove the limitation when this limitation is removed
+    #"protobuf<=3.20", # Google bigquery client require protobuf <= 3.20.0. We can remove the limitation when this limitation is removed
     "openlineage-airflow>=0.17.0",
     "apache-airflow-providers-microsoft-azure",
     "azure-storage-blob",

python-sdk/src/astro/sql/operators/cleanup.py

@@ -225,6 +225,8 @@ def _get_executor_from_job_id(job_id: int) -> str | None:
 
         with create_session() as session:
             job = session.get(Job, job_id)
+        if job.executor_class is None and job.executor:
+            return type(job.executor).__name__
         return job.executor_class if job else None
 
     def get_all_task_outputs(self, context: Context) -> list[BaseTable]:

python-sdk/tests/modified_constraint_file.txt → python-sdk/tests/modified_constraint_file_2.2.5.txt

@@ -36,7 +36,7 @@ PyGithub==1.54.1
 PyHive==0.6.5
 PyJWT==1.7.1
 PyNaCl==1.5.0
-PyYAML==5.4.1
+PyYAML==5.3.1
 Pygments==2.11.2
 Rx==3.2.0
 SQLAlchemy-JSONField==1.0.0