SNOW-3392625: Add artifact_repository support to dbapi udtf_configs (#4184)

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

@@ -2350,6 +2350,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/datasource/test_oracledb.py

@@ -416,3 +416,47 @@ def create_connection_with_params(user=None, password=None, dsn=None, **kwargs):
         udtf_configs=udtf_configs,
     )
     Utils.check_answer(df, oracledb_real_data)
+
+
+@pytest.mark.udf
+def test_udtf_ingestion_oracledb_with_artifact_repository(session):
+    """Verify artifact_repository in udtf_configs is forwarded to the UDTF creation DDL."""
+    from tests.parameters import ORACLEDB_CONNECTION_PARAMETERS
+
+    his = session.query_history()
+
+    def create_connection_oracledb():
+        import oracledb
+
+        host = ORACLEDB_CONNECTION_PARAMETERS["host"]
+        port = ORACLEDB_CONNECTION_PARAMETERS["port"]
+        service_name = ORACLEDB_CONNECTION_PARAMETERS["service_name"]
+        username = ORACLEDB_CONNECTION_PARAMETERS["username"]
+        password = ORACLEDB_CONNECTION_PARAMETERS["password"]
+        dsn = f"{host}:{port}/{service_name}"
+        connection = oracledb.connect(user=username, password=password, dsn=dsn)
+        return connection
+
+    df = session.read.dbapi(
+        create_connection_oracledb,
+        table="ALL_TYPE_TABLE",
+        udtf_configs={
+            "external_access_integration": ORACLEDB_TEST_EXTERNAL_ACCESS_INTEGRATION,
+            "artifact_repository": "SNOWPARK_PYTHON_TEST_REPOSITORY",
+            "packages": ["oracledb", "snowflake-snowpark-python", "cloudpickle"],
+        },
+    ).order_by("ID")
+
+    Utils.check_answer(df, oracledb_real_data)
+    assert df.schema == oracledb_real_schema
+
+    # check that the UDTF creation DDL includes the artifact repository
+    found = any(
+        (
+            "CREATE" in q.sql_text.upper()
+            and "FUNCTION" in q.sql_text.upper()
+            and "SNOWPARK_PYTHON_TEST_REPOSITORY" in q.sql_text
+        )
+        for q in his.queries
+    )
+    assert found, "artifact_repository not found in UDTF creation DDL"