SNOW-3375255: enable Python 3.14 UDF tests (#4172)

Author: sfc-gh-bkogan

.github/workflows/daily_precommit.yml

@@ -614,7 +614,7 @@ jobs:
       - name: Install tox
         run: uv pip install tox --system
       # TODO: enable doctest for 3.14
-      - if: ${{ contains('macos', matrix.os.download_name) && matrix.python-version != '3.14' }}
+      - if: ${{ contains('macos', matrix.os.download_name) }}
         name: Run doctests
         run: python -m tox -e "py${PYTHON_VERSION}-doctest-notudf-ci"
         env:
@@ -623,9 +623,7 @@ jobs:
           PYTEST_ADDOPTS: --color=yes --tb=short --disable_cte_optimization
           TOX_PARALLEL_NO_SPINNER: 1
         shell: bash
-      # TODO: enable for 3.14
-      - if: ${{ matrix.python-version != '3.14' }}
-        name: Run tests (excluding doctests)
+      - name: Run tests (excluding doctests)
         run: python -m tox -e "py${PYTHON_VERSION/\./}-dailynotdoctest-ci"
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
@@ -635,18 +633,6 @@ jobs:
           SNOWPARK_PYTHON_API_TEST_BUCKET_PATH: ${{ secrets.SNOWPARK_PYTHON_API_TEST_BUCKET_PATH }}
           SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION: ${{ vars.SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION }}
         shell: bash
-      # TODO: remove the test below and run udf tests for 3.14
-      - if: ${{ matrix.python-version == '3.14' }}
-        name: Run tests (excluding udf, doctests)
-        run: python -m tox -e "py${PYTHON_VERSION/\./}-dailynotdoctestnotudf-ci"
-        env:
-          PYTHON_VERSION: ${{ matrix.python-version }}
-          cloud_provider: ${{ matrix.cloud-provider }}
-          PYTEST_ADDOPTS: --color=yes --tb=short --disable_cte_optimization
-          TOX_PARALLEL_NO_SPINNER: 1
-          SNOWPARK_PYTHON_API_TEST_BUCKET_PATH: ${{ secrets.SNOWPARK_PYTHON_API_TEST_BUCKET_PATH }}
-          SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION: ${{ vars.SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION }}
-        shell: bash
       - name: Combine coverages
         run: python -m tox -e coverage --skip-missing-interpreters false
         shell: bash

.github/workflows/precommit.yml

@@ -175,8 +175,7 @@ jobs:
       - name: Install tox
         run: uv pip install tox --system
       # we only run doctest on macos
-      # TODO: enable doctest for 3.14
-      - if: ${{ matrix.os == 'macos-latest' && matrix.python-version != '3.14' }}
+      - if: ${{ matrix.os == 'macos-latest' }}
         name: Run doctests
         run: python -m tox -e "py${PYTHON_VERSION}-doctest-notudf-ci"
         env:
@@ -188,7 +187,7 @@ jobs:
           # For example, see https://github.com/snowflakedb/snowpark-python/pull/681
         shell: bash
       # do not run other tests for macos
-      - if: ${{ matrix.os != 'macos-latest' && matrix.python-version != '3.14' }}
+      - if: ${{ matrix.os != 'macos-latest' }}
         name: Run tests (excluding doctests)
         run: python -m tox -e "py${PYTHON_VERSION/\./}-notdoctest-ci"
         env:
@@ -199,19 +198,6 @@ jobs:
           SNOWPARK_PYTHON_API_TEST_BUCKET_PATH: ${{ secrets.SNOWPARK_PYTHON_API_TEST_BUCKET_PATH }}
           SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION: ${{ vars.SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION }}
         shell: bash
-      # TODO: Remove the test below and run udf tests for 3.14
-      # for 3.14, skip udf, doctest
-      - if: ${{ matrix.os != 'macos-latest' && matrix.python-version == '3.14' }}
-        name: Run tests (excluding udf, doctests)
-        run: python -m tox -e "py${PYTHON_VERSION/\./}-notudfdoctest-ci"
-        env:
-          PYTHON_VERSION: ${{ matrix.python-version }}
-          cloud_provider: ${{ matrix.cloud-provider }}
-          PYTEST_ADDOPTS: --color=yes --tb=short
-          TOX_PARALLEL_NO_SPINNER: 1
-          SNOWPARK_PYTHON_API_TEST_BUCKET_PATH: ${{ secrets.SNOWPARK_PYTHON_API_TEST_BUCKET_PATH }}
-          SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION: ${{ vars.SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION }}
-        shell: bash
       - name: Install MS ODBC Driver (Ubuntu only)
         if: ${{ contains(matrix.os, 'ubuntu') }}
         run: |
@@ -222,8 +208,7 @@ jobs:
         shell: bash
       - name: Run data source tests
         # psycopg2 is not supported on macos 3.9
-        # TODO: enable datasource tests for 3.14
-        if: ${{ !(matrix.os == 'macos-latest' && matrix.python-version == '3.9') && !(matrix.python-version == '3.14') }}
+        if: ${{ !(matrix.os == 'macos-latest' && matrix.python-version == '3.9') }}
         run: python -m tox -e datasource
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}

src/conftest.py

@@ -13,6 +13,7 @@
 
 from snowflake.snowpark import Session
 from snowflake.snowpark.functions import to_timestamp
+from snowflake.snowpark.context import _DEFAULT_ARTIFACT_REPOSITORY
 
 logging.getLogger("snowflake.connector").setLevel(logging.ERROR)
 
@@ -72,6 +73,9 @@ def add_snowpark_session(doctest_namespace, pytestconfig):
             session.sql(
                 f"GRANT ALL PRIVILEGES ON SCHEMA {TEST_SCHEMA} TO ROLE PUBLIC"
             ).collect()
+            session.sql(
+                f"ALTER SCHEMA SET DEFAULT_PYTHON_ARTIFACT_REPOSITORY = {_DEFAULT_ARTIFACT_REPOSITORY}"
+            ).collect()
             session.use_schema(TEST_SCHEMA)
         doctest_namespace["session"] = session
         yield

src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py

@@ -167,7 +167,14 @@ def udtf_ingestion(
         statement_params: Optional[Dict[str, str]] = None,
         _emit_ast: bool = True,
     ) -> "snowflake.snowpark.DataFrame":
-        from snowflake.snowpark._internal.data_source.utils import UDTF_PACKAGE_MAP
+        from snowflake.snowpark._internal.data_source.utils import (
+            resolve_udtf_packages,
+        )
+
+        resolved_packages = packages or resolve_udtf_packages(
+            self.dbms_type,
+            artifact_repository or session._get_default_artifact_repository(),
+        )
 
         udtf_name = random_name_for_temp_object(TempObjectType.FUNCTION)
         with measure_time() as udtf_register_time:
@@ -186,7 +193,7 @@ def udtf_ingestion(
                     ]
                 ),
                 external_access_integrations=[external_access_integrations],
-                packages=packages or UDTF_PACKAGE_MAP.get(self.dbms_type),
+                packages=resolved_packages,
                 imports=imports,
                 artifact_repository=artifact_repository,
                 statement_params=statement_params,

src/snowflake/snowpark/_internal/data_source/utils.py

@@ -34,7 +34,11 @@
 from snowflake.snowpark._internal.data_source import DataSourceReader
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.snowpark._internal.utils import get_temp_type_for_object
-from snowflake.snowpark.exceptions import SnowparkDataframeReaderException
+from snowflake.snowpark.context import _PYPI_SHARED_REPOSITORY
+from snowflake.snowpark.exceptions import (
+    SnowparkClientException,
+    SnowparkDataframeReaderException,
+)
 from snowflake.snowpark.types import StructType
 
 from typing import TYPE_CHECKING
@@ -98,7 +102,11 @@ class DRIVER_TYPE(str, Enum):
     DRIVER_TYPE.PYMYSQL: PymysqlDriver,
 }
 
-UDTF_PACKAGE_MAP = {
+# Default UDTF package list, suitable for Snowflake's Anaconda shared
+# repository. The Snowflake Anaconda channel ships conda builds of these
+# packages with the necessary native libraries (e.g., libpq for psycopg2,
+# msodbcsql for pyodbc) bundled, so the source distribution names work.
+_ANACONDA_UDTF_PACKAGE_MAP = {
     DBMS_TYPE.ORACLE_DB: ["oracledb>=2.0.0,<4.0.0", "snowflake-snowpark-python"],
     DBMS_TYPE.SQLITE_DB: ["snowflake-snowpark-python"],
     DBMS_TYPE.SQL_SERVER_DB: [
@@ -114,6 +122,67 @@ class DRIVER_TYPE(str, Enum):
     DBMS_TYPE.MYSQL_DB: ["pymysql>=1.0.0,<2.0.0", "snowflake-snowpark-python"],
 }
 
+# UDTF package list when using the PyPI shared repository. The server-side
+# UDTF install sandbox refuses to compile source distributions (sdists), so
+# every package here must be wheel-installable from PyPI. Differences from
+# the Anaconda map:
+#   - Postgres uses ``psycopg2-binary`` because ``psycopg2`` on PyPI is
+#     sdist-only; ``psycopg2-binary`` is the wheel-packaged equivalent.
+#   - SQL Server has no PyPI-installable equivalent of ``msodbcsql`` (it is
+#     Microsoft's ODBC driver, distributed as a system package), so the
+#     UDTF path cannot work on PyPI today.
+#   - Databricks depends on ``databricks-sql-connector``, which transitively
+#     requires ``thrift``; ``thrift`` on PyPI is sdist-only for every
+#     version, so the server cannot install it.
+# These PyPI gaps are independent of the Python version; they apply to any
+# session whose default artifact repository is PyPI (most commonly Python
+# 3.14+, where PyPI is the global default).
+_PYPI_UDTF_PACKAGE_MAP = {
+    DBMS_TYPE.ORACLE_DB: ["oracledb>=2.0.0,<4.0.0", "snowflake-snowpark-python"],
+    DBMS_TYPE.SQLITE_DB: ["snowflake-snowpark-python"],
+    DBMS_TYPE.POSTGRES_DB: [
+        "psycopg2-binary>=2.0.0,<3.0.0",
+        "snowflake-snowpark-python",
+    ],
+    DBMS_TYPE.MYSQL_DB: ["pymysql>=1.0.0,<2.0.0", "snowflake-snowpark-python"],
+    # SQL_SERVER_DB and DATABRICKS_DB intentionally omitted - see
+    # resolve_udtf_packages for the user-facing error.
+}
+
+# Backwards-compatible alias for callers (and external code) that imported
+# the old map. New code should use :func:`resolve_udtf_packages`.
+UDTF_PACKAGE_MAP = _ANACONDA_UDTF_PACKAGE_MAP
+
+
+def resolve_udtf_packages(
+    dbms_type: "DBMS_TYPE", artifact_repository: Optional[str]
+) -> List[str]:
+    """Return the default UDTF package list for ``dbms_type``.
+
+    Picks the package list appropriate for ``artifact_repository``. When the
+    repository is the PyPI shared repository, some DBMSes have no working
+    package set (their dependencies are not wheel-installable from PyPI, and
+    the server-side UDTF sandbox refuses to compile sdists); this raises
+    :class:`SnowparkClientException` with guidance to switch repositories.
+    """
+    if artifact_repository == _PYPI_SHARED_REPOSITORY:
+        packages = _PYPI_UDTF_PACKAGE_MAP.get(dbms_type)
+        if packages is None:
+            raise SnowparkClientException(
+                f"DataFrameReader.dbapi server-side UDTF ingestion for "
+                f"{dbms_type.value} is not supported when the session's "
+                f"default artifact repository is PyPI: the required "
+                f"packages are not wheel-installable from PyPI, and the "
+                f"server-side UDTF install sandbox refuses to compile "
+                f"source distributions. Switch to the Anaconda artifact "
+                f"repository (on Python 3.14+, PyPI is the client-side "
+                f"default; on older Python versions, Anaconda is the "
+                f"default but may have been overridden at the account, "
+                f"database, or schema level)."
+            )
+        return packages
+    return _ANACONDA_UDTF_PACKAGE_MAP.get(dbms_type)
+
 
 def get_jdbc_dbms(jdbc_url: str) -> str:
     """

src/snowflake/snowpark/_internal/udf_utils.py

@@ -1247,7 +1247,7 @@ def resolve_imports_and_packages(
     if artifact_repository != _ANACONDA_SHARED_REPOSITORY:
         # Non-conda artifact repository - skip conda-based package resolution
         resolved_packages = []
-        if not packages and session:
+        if packages is None and session:
             resolved_packages = list(
                 session._resolve_packages(
                     [],
@@ -1256,7 +1256,7 @@ def resolve_imports_and_packages(
                     include_pandas=is_pandas_udf,
                 )
             )
-        elif packages:
+        elif packages is not None:
             if not all(isinstance(package, str) for package in packages):
                 raise TypeError(
                     "Non-conda artifact repository requires that all packages be passed as str."
@@ -1615,7 +1615,6 @@ def create_python_udf_or_sp(
         is_ddl_on_temp_object=not is_permanent,
         statement_params=statement_params,
     )
-
     if comment is not None:
         object_signature_sql = f"{object_name}({','.join(input_sql_types)})"
         comment = escape_single_quotes(comment)

src/snowflake/snowpark/context.py

@@ -5,6 +5,7 @@
 
 """Context module for Snowpark."""
 import logging
+import sys
 from typing import Callable, Optional
 
 import snowflake.snowpark
@@ -168,8 +169,14 @@
 
 # The fully qualified name of the Anaconda shared repository (conda channel).
 _ANACONDA_SHARED_REPOSITORY = "snowflake.snowpark.anaconda_shared_repository"
-# In case of failures or the current default artifact repository is unset, we fallback to this
-_DEFAULT_ARTIFACT_REPOSITORY = _ANACONDA_SHARED_REPOSITORY
+# The fully qualified name of the PyPI shared repository (pypi channel).
+_PYPI_SHARED_REPOSITORY = "snowflake.snowpark.pypi_shared_repository"
+# In case of failures and for routing to the right session package store, we use this
+_DEFAULT_ARTIFACT_REPOSITORY = (
+    _ANACONDA_SHARED_REPOSITORY
+    if sys.version_info < (3, 14)
+    else _PYPI_SHARED_REPOSITORY
+)
 
 
 def configure_development_features(

tests/integ/conftest.py

@@ -4,20 +4,21 @@
 #
 import os
 from logging import getLogger
-import sys
 from typing import Dict
 
 import pytest
 
 import snowflake.connector
 from snowflake.snowpark import Session
+from snowflake.snowpark.context import _DEFAULT_ARTIFACT_REPOSITORY
 from snowflake.snowpark._internal.utils import set_ast_state, AstFlagSource
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.mock._connection import MockServerConnection
 from tests.ast.ast_test_utils import (
     close_full_ast_validation_mode,
     setup_full_ast_validation_mode,
 )
+from tests.integ.session_parameters import set_up_test_session_parameters
 from tests.parameters import CONNECTION_PARAMETERS
 from tests.utils import (
     TEST_SCHEMA,
@@ -286,6 +287,9 @@ def test_schema(connection, local_testing_mode) -> None:
             cursor.execute(
                 f"GRANT ALL PRIVILEGES ON SCHEMA {TEST_SCHEMA} TO ROLE PUBLIC"
             )
+            cursor.execute(
+                f"ALTER SCHEMA SET DEFAULT_PYTHON_ARTIFACT_REPOSITORY = {_DEFAULT_ARTIFACT_REPOSITORY}"
+            )
             yield
             cursor.execute(f"DROP SCHEMA IF EXISTS {TEST_SCHEMA}")
 
@@ -355,13 +359,7 @@ def session(
         )
 
     # TODO: SNOW-2346239: Set parameter on user level instead of in config file
-    if not local_testing_mode:
-        session.sql(
-            "alter session set ENABLE_EXTRACTION_PUSHDOWN_EXTERNAL_PARQUET_FOR_COPY_PHASE_I='Track';"
-        ).collect()
-        session.sql("alter session set ENABLE_ROW_ACCESS_POLICY=true").collect()
-        if sys.version_info.major == 3 and sys.version_info.minor == 14:
-            session.sql("alter session set ENABLE_PYTHON_3_14=true").collect()
+    set_up_test_session_parameters(session, local_testing_mode)
 
     try:
         yield session
@@ -412,6 +410,7 @@ def profiler_session(
             integration2,
             integration3,
         )
+    set_up_test_session_parameters(session, local_testing_mode)
     try:
         yield session
     finally:
@@ -434,6 +433,9 @@ def temp_schema(connection, session, local_testing_mode) -> None:
             cursor.execute(
                 f"GRANT ALL PRIVILEGES ON SCHEMA {temp_schema_name} TO ROLE PUBLIC"
             )
+            cursor.execute(
+                f"ALTER SCHEMA SET DEFAULT_PYTHON_ARTIFACT_REPOSITORY = {_DEFAULT_ARTIFACT_REPOSITORY}"
+            )
             yield temp_schema_name
             cursor.execute(f"DROP SCHEMA IF EXISTS {temp_schema_name}")
 

tests/integ/datasource/test_databricks.py

@@ -42,7 +42,7 @@
     databricks_unicode_schema,
     databricks_double_quoted_schema,
 )
-from tests.utils import IS_IN_STORED_PROC, IS_MACOS, Utils
+from tests.utils import IS_IN_STORED_PROC, IS_MACOS, Utils, IS_PY314
 
 DEPENDENCIES_PACKAGE_UNAVAILABLE = True
 try:
@@ -177,6 +177,10 @@ def test_double_quoted_column_databricks(session, custom_schema):
     [("table", TEST_TABLE_NAME), ("query", f"(SELECT * FROM {TEST_TABLE_NAME})")],
 )
 @pytest.mark.udf
+@pytest.mark.skipif(
+    IS_PY314,
+    reason="databricks-sql-connector's thrift dependency has no Python 3.14 wheel on PyPI; server-side UDTF install fails on the default repo.",
+)
 def test_udtf_ingestion_databricks(session, input_type, input_value, caplog):
     # we define here to avoid test_databricks.py to be pickled and unpickled in UDTF
     def local_create_databricks_connection():
@@ -286,6 +290,10 @@ def test_session_init(session):
         )
 
 
+@pytest.mark.skipif(
+    IS_PY314,
+    reason="databricks-sql-connector's thrift dependency has no Python 3.14 wheel on PyPI; server-side UDTF install fails on the default repo.",
+)
 def test_session_init_udtf(session):
     udtf_configs = {
         "external_access_integration": DATABRICKS_TEST_EXTERNAL_ACCESS_INTEGRATION