SNOW-940615: fix pyarrow import if pyarrow iterator is not compiled (#1769)

Author: sfc-gh-aling

.github/workflows/build_test.yml

@@ -210,6 +210,38 @@ jobs:
           PYTEST_ADDOPTS: --color=yes --tb=short
         shell: bash
 
+  test-noarrowextension:
+    name: No Arrow Extension Test ${{ matrix.os.download_name }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}
+    needs: lint
+    runs-on: ${{ matrix.os.image_name }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+         - image_name: ubuntu-latest
+           download_name: linux
+        python-version: [3.8]
+        cloud-provider: [aws]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Upgrade setuptools, pip and wheel
+        run: python -m pip install -U setuptools pip wheel
+      - name: Install tox
+        run: python -m pip install tox
+      - name: Run tests
+        run: python -m tox -e noarrowextension
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short
+        shell: bash
+
   test-fips:
     name: Test FIPS linux-3.8-${{ matrix.cloud-provider }}
     needs: build

DESCRIPTION.md

@@ -7,12 +7,11 @@ https://docs.snowflake.com/
 Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python
 
 # Release Notes
-- v3.2.2(TBD)
-  - Fixed issue with connection diagnostics failing to complete certificate checks.
-
 - v3.3.1(Unreleased)
 
   - Added for non-Windows platforms command suggestions (chown/chmod) for insufficient file permissions of config files.
+  - Fixed issue with connection diagnostics failing to complete certificate checks.
+  - Fixed issue that arrow iterator causes `ImportError` when the c extensions are not compiled.
 
 - v3.3.0(October 10,2023)
 

setup.py

@@ -42,6 +42,10 @@
 extensions = None
 cmd_class = {}
 
+SNOWFLAKE_DISABLE_COMPILE_ARROW_EXTENSIONS = os.environ.get(
+    "SNOWFLAKE_DISABLE_COMPILE_ARROW_EXTENSIONS", "false"
+).lower() in ("y", "yes", "t", "true", "1", "on")
+
 try:
     import numpy
     import pyarrow
@@ -56,7 +60,7 @@
     )
     _ABLE_TO_COMPILE_EXTENSIONS = False
 
-if _ABLE_TO_COMPILE_EXTENSIONS:
+if _ABLE_TO_COMPILE_EXTENSIONS and not SNOWFLAKE_DISABLE_COMPILE_ARROW_EXTENSIONS:
     pyarrow_version = tuple(int(x) for x in pyarrow.__version__.split("."))
     extensions = cythonize(
         [

src/snowflake/connector/result_batch.py

@@ -14,13 +14,10 @@
 from typing import TYPE_CHECKING, Any, Callable, Iterator, NamedTuple, Sequence
 
 from .arrow_context import ArrowConverterContext
-from .arrow_iterator import PyArrowIterator
 from .compat import OK, UNAUTHORIZED, urlparse
 from .constants import FIELD_TYPES, IterUnit
 from .errorcode import ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE, ER_NO_PYARROW
 from .errors import Error, InterfaceError, NotSupportedError, ProgrammingError
-from .nanoarrow_arrow_iterator import PyArrowRowIterator as NanoarrowRowIterator
-from .nanoarrow_arrow_iterator import PyArrowTableIterator as NanoarrowTableIterator
 from .network import (
     RetryRequest,
     get_http_retryable_error,
@@ -65,7 +62,9 @@ def _create_vendored_arrow_iterator(
     numpy: bool,
     number_to_decimal: bool,
     row_unit: IterUnit,
-) -> PyArrowIterator:
+):
+    from .arrow_iterator import PyArrowIterator
+
     logger.debug("Using vendored arrow as the arrow data converter")
     iter = PyArrowIterator(
         None,
@@ -87,10 +86,12 @@ def _create_nanoarrow_iterator(
     numpy: bool,
     number_to_decimal: bool,
     row_unit: IterUnit,
-) -> NanoarrowRowIterator | NanoarrowTableIterator:
+):
+    from .nanoarrow_arrow_iterator import PyArrowRowIterator, PyArrowTableIterator
+
     logger.debug("Using nanoarrow as the arrow data converter")
     return (
-        NanoarrowRowIterator(
+        PyArrowRowIterator(
             None,
             data,
             context,
@@ -99,7 +100,7 @@ def _create_nanoarrow_iterator(
             number_to_decimal,
         )
         if row_unit == IterUnit.ROW_UNIT
-        else NanoarrowTableIterator(
+        else PyArrowTableIterator(
             None,
             data,
             context,

test/unit/test_connection_diagnostic.py

@@ -5,8 +5,20 @@
 
 from __future__ import annotations
 
+import pytest
+
+pytestmark = pytest.mark.skipolddriver  # old test driver tests won't run this module
+
+
 import snowflake.connector
-from snowflake.connector.connection_diagnostic import ConnectionDiagnostic, _decode_dict
+
+try:
+    from snowflake.connector.connection_diagnostic import (
+        ConnectionDiagnostic,
+        _decode_dict,
+    )
+except ImportError:
+    pass
 
 
 def test_https_host_report(caplog):

tox.ini

@@ -96,6 +96,15 @@ passenv = {[testenv]passenv}
 commands =
     {env:SNOWFLAKE_PYTEST_CMD} -m "not skipolddriver" -vvv {posargs:} test
 
+[testenv:noarrowextension]
+basepython = python3.8
+skip_install = True
+description = run import with no arrow extension under {basepython}
+setenv = SNOWFLAKE_DISABLE_COMPILE_ARROW_EXTENSIONS=1
+commands =
+    pip install .
+    python -c 'import snowflake.connector.result_batch'
+
 [testenv:coverage]
 description = [run locally after tests]: combine coverage data and create report
 ;              generates a diff coverage against origin/master (can be changed by setting DIFF_AGAINST env var)