SNOW-3392625: Add artifact_repository support to dbapi udtf_configs

Forward the artifact_repository parameter from dbapi() udtf_configs
through to session.udtf.register(), enabling users to specify a
custom artifact repository (e.g. PyPI) for packages used by the
internal UDTF created during distributed dbapi ingestion.

Changes:
- dataframe_reader.py: extract artifact_repository from udtf_configs
- datasource_partitioner.py: forward through _udtf_ingestion wrapper
- base_driver.py: accept and forward to session.udtf.register()
- CHANGELOG.md: add new feature entry
- test_data_source_api.py: add test_dbapi_udtf_artifact_repository

.... Generated with [Cortex Code](https://docs.snowflake.com/en/user-guide/cortex-code/cortex-code)

Co-Authored-By: Cortex Code <noreply@snowflake.com>

Author: sfc-gh-qding

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### Snowpark Python API Updates
 
+#### New Features
+
+- Added `artifact_repository` support to `udtf_configs` in `session.read.dbapi()`, enabling users to specify a custom artifact repository (e.g. PyPI) for packages used by the internal UDTF during distributed ingestion.
+
 #### Bug Fixes
 
 - Fixed a bug where `TRY_CAST` reader option is ignored when calling `DataFrameReader.schema().csv()`.

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

@@ -184,6 +184,7 @@ def _udtf_ingestion(
         fetch_size: int = 1000,
         imports: Optional[List[str]] = None,
         packages: Optional[List[str]] = None,
+        artifact_repository: Optional[str] = None,
         session_init_statement: Optional[List[str]] = None,
         query_timeout: Optional[int] = 0,
         statement_params: Optional[Dict[str, str]] = None,
@@ -197,6 +198,7 @@ def _udtf_ingestion(
             fetch_size,
             imports,
             packages,
+            artifact_repository,
             session_init_statement,
             query_timeout,
             statement_params,

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

@@ -161,6 +161,7 @@ def udtf_ingestion(
         fetch_size: int = 1000,
         imports: Optional[List[str]] = None,
         packages: Optional[List[str]] = None,
+        artifact_repository: Optional[str] = None,
         session_init_statement: Optional[List[str]] = None,
         query_timeout: Optional[int] = 0,
         statement_params: Optional[Dict[str, str]] = None,
@@ -187,6 +188,7 @@ def udtf_ingestion(
                 external_access_integrations=[external_access_integrations],
                 packages=packages or UDTF_PACKAGE_MAP.get(self.dbms_type),
                 imports=imports,
+                artifact_repository=artifact_repository,
                 statement_params=statement_params,
                 _emit_ast=_emit_ast,  # internal function call, _emit_ast will be set to False by the caller
             )

src/snowflake/snowpark/dataframe_reader.py

@@ -2305,6 +2305,7 @@ def create_oracledb_connection():
                 fetch_size=fetch_size,
                 imports=udtf_configs.get("imports", None),
                 packages=udtf_configs.get("packages", None),
+                artifact_repository=udtf_configs.get("artifact_repository", None),
                 session_init_statement=session_init_statement,
                 query_timeout=query_timeout,
                 statement_params=statements_params_for_telemetry,

tests/integ/test_data_source_api.py

@@ -79,7 +79,7 @@
     sql_server_create_connection_unicode_data,
     sql_server_create_connection_double_quoted_data,
 )
-from tests.utils import Utils, IS_WINDOWS, RUNNING_ON_JENKINS
+from tests.utils import Utils, IS_WINDOWS, RUNNING_ON_JENKINS, IS_NOT_ON_GITHUB
 
 try:
     import pandas  # noqa: F401
@@ -1924,3 +1924,66 @@ def create_mock_oracledb_v34_or_higher():
     dbms_type, driver_type = detect_dbms(create_mock_oracledb_v34_or_higher())
     assert dbms_type == DBMS_TYPE.ORACLE_DB
     assert driver_type == DRIVER_TYPE.ORACLEDB
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="artifact repository not supported in local testing",
+)
+@pytest.mark.skipif(IS_NOT_ON_GITHUB, reason="need resources")
+@pytest.mark.skipif(
+    RUNNING_ON_JENKINS,
+    reason="SNOW-2089683: oracledb real connection test failed on jenkins",
+)
+@pytest.mark.udf
+def test_dbapi_udtf_artifact_repository(session, ast_enabled):
+    """Verify artifact_repository in udtf_configs is forwarded to the UDTF."""
+    if ast_enabled:
+        pytest.skip("TODO: dbapi has not implemented ast yet, skip the test for now")
+
+    import sys
+
+    if sys.version_info < (3, 9):
+        pytest.skip("artifact repository requires Python 3.9+")
+
+    def create_connection():
+        class FakeConnection:
+            def cursor(self):
+                class FakeCursor:
+                    def execute(self, query):
+                        pass
+
+                    @property
+                    def description(self):
+                        return [("c1", int, None, None, None, None, None)]
+
+                    def fetchmany(self, *args, **kwargs):
+                        return None
+
+                return FakeCursor()
+
+        return FakeConnection()
+
+    his = session.query_history()
+    df = session.read.dbapi(
+        create_connection,
+        table="Fake",
+        custom_schema="c1 INT",
+        udtf_configs={
+            "external_access_integration": ORACLEDB_TEST_EXTERNAL_ACCESS_INTEGRATION,
+            "packages": ["snowflake-snowpark-python", "cloudpickle"],
+            "artifact_repository": "SNOWPARK_PYTHON_TEST_REPOSITORY",
+        },
+    )
+    df.select("*").collect()
+
+    # Verify the CREATE FUNCTION DDL includes the artifact repository
+    create_function_queries = [
+        q.sql_text
+        for q in his.queries
+        if "CREATE" in q.sql_text.upper() and "FUNCTION" in q.sql_text.upper()
+    ]
+    assert len(create_function_queries) > 0, "Expected at least one CREATE FUNCTION query"
+    assert any(
+        "SNOWPARK_PYTHON_TEST_REPOSITORY" in q for q in create_function_queries
+    ), f"artifact_repository not found in UDTF DDL: {create_function_queries}"