fix(snowflake): support snowflake 3.3.0 (nanoarrow)

ibis-project · Oct 13, 2023 · a0f24e8 · a0f24e8
1 parent ca940d9
commit a0f24e8
Show file tree

Hide file tree

Showing 6 changed files with 606 additions and 599 deletions.
diff --git a/.github/workflows/ibis-backends-cloud.yml b/.github/workflows/ibis-backends-cloud.yml
@@ -42,18 +42,8 @@ jobs:
         backend:
           - name: bigquery
             title: BigQuery
-            deps: []
           - name: snowflake
             title: Snowflake
-            deps: []
-        include:
-          - python-version: "3.10"
-            backend:
-              name: snowflake
-              title: Snowflake with nanoarrow
-              deps:
-                - snowflake-connector-python[pandas]@3.3.0b1
-                - --optional
     steps:
       - name: checkout
         uses: actions/checkout@v4
@@ -65,20 +55,14 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - uses: syphar/restore-virtualenv@v1
-        if: toJSON(matrix.backend.deps) == '[]'
         with:
           requirement_files: poetry.lock
           custom_cache_key_element: ${{ matrix.backend.name }}-${{ steps.install_python.outputs.python-version }}
 
       - name: upgrade pip and install poetry
         run: python -m pip install --upgrade pip 'poetry==1.6.1'
 
-      - name: install a prerelease version of `snowflake-connector-python`
-        if: toJSON(matrix.backend.deps) != '[]'
-        run: poetry add ${{ join(matrix.backend.deps, ' ') }}
-
       - uses: syphar/restore-pip-download-cache@v1
-        if: toJSON(matrix.backend.deps) == '[]'
         with:
           requirement_files: poetry.lock
           custom_cache_key_element: ${{ matrix.backend.name }}-${{ steps.install_python.outputs.python-version }}

diff --git a/.github/workflows/ibis-backends.yml b/.github/workflows/ibis-backends.yml
@@ -38,52 +38,9 @@ env:
   FORCE_COLOR: "1"
 
 jobs:
-  gen_lockfile_backends:
-    name: Generate Poetry Lockfile for non-Snowflake Backends
-    runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version:
-          - "3.9"
-          - "3.11"
-    steps:
-      - name: checkout
-        uses: actions/checkout@v4
-
-      - name: install python
-        id: install_python
-        uses: actions/setup-python@v4
-        with:
-          python-version: "${{ matrix.python-version }}"
-
-      - name: install poetry
-        run: python -m pip install --upgrade pip 'poetry==1.6.1'
-
-      - uses: syphar/restore-pip-download-cache@v1
-        with:
-          requirement_files: poetry.lock
-          custom_cache_key_element: ${{ steps.install_python.outputs.python-version }}
-
-      - name: remove snowflake deps that are not compatible with everything else
-        run: poetry remove snowflake-sqlalchemy snowflake-connector-python
-
-      - name: update deps originally constrained by snowflake
-        run: poetry update numpy pandas pyarrow datafusion
-
-      - name: upload deps file
-        uses: actions/upload-artifact@v3
-        with:
-          name: backend-deps-${{ matrix.python-version }}
-          path: |
-            pyproject.toml
-            poetry.lock
-
   test_backends:
     name: ${{ matrix.backend.title }} ${{ matrix.os }} python-${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
-    needs:
-      - gen_lockfile_backends
     env:
       SQLALCHEMY_WARN_20: "1"
     strategy:
@@ -350,20 +307,6 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: download poetry lockfile
-        uses: actions/download-artifact@v3
-        with:
-          name: backend-deps-${{ matrix.python-version }}
-          path: deps
-
-      - name: pull out lockfile
-        shell: bash
-        run: |
-          set -euo pipefail
-
-          mv -f deps/* .
-          rm -r deps
-
       - uses: syphar/restore-pip-download-cache@v1
         with:
           requirement_files: poetry.lock
@@ -684,7 +627,8 @@ jobs:
   test_backends_sqlalchemy2:
     name: SQLAlchemy 2 ${{ matrix.backend.title }} ${{ matrix.os }} python-${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
-    needs: gen_lockfile_sqlalchemy2
+    needs:
+      - gen_lockfile_sqlalchemy2
     strategy:
       fail-fast: false
       matrix:

diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
@@ -3,6 +3,7 @@
 import contextlib
 import functools
 import glob
+import importlib
 import inspect
 import itertools
 import json
@@ -19,7 +20,7 @@
 
 import pyarrow as pa
 import sqlalchemy as sa
-from snowflake.sqlalchemy import ARRAY, DOUBLE, OBJECT, URL
+from packaging.version import parse as vparse
 from sqlalchemy.ext.compiler import compiles
 
 import ibis
@@ -35,9 +36,21 @@
     AlchemyCrossSchemaBackend,
     AlchemyExprTranslator,
 )
-from ibis.backends.snowflake.converter import SnowflakePandasData
-from ibis.backends.snowflake.datatypes import SnowflakeType
-from ibis.backends.snowflake.registry import operation_registry
+
+with warnings.catch_warnings():
+    if vparse(importlib.metadata.version("snowflake-connector-python")) >= vparse(
+        "3.3.0"
+    ):
+        warnings.filterwarnings(
+            "ignore",
+            message="You have an incompatible version of 'pyarrow' installed",
+            category=UserWarning,
+        )
+    from snowflake.sqlalchemy import ARRAY, DOUBLE, OBJECT, URL
+
+    from ibis.backends.snowflake.converter import SnowflakePandasData
+    from ibis.backends.snowflake.datatypes import SnowflakeType
+    from ibis.backends.snowflake.registry import operation_registry
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Mapping