Merge branch 'main' into aalam-SNOW-2257191-cte-join-bugfix

Author: sfc-gh-aling

CHANGELOG.md

@@ -56,13 +56,8 @@
 #### Improvements
 
 - Enhanced `DataFrame.sort()` to support `ORDER BY ALL` when no columns are specified.
-- Catalog API now uses SQL commands instead of SnowAPI calls. This new implementation is more reliable now.
 - Removed experimental warning from `Session.cte_optimization_enabled`.
 
-#### Dependency Updates
-
-- Catalog API no longer uses types declared in `snowflake.core` and therefore this dependency was removed.
-
 #### Bug Fixes
 
 - Fixed with a bug when sql generation when joining two `DataFrame`s created using `DataFrame.alias` and CTE optimization is enabled.

setup.py

@@ -63,6 +63,7 @@
     "pytest-assume",  # sql counter check
     "decorator",  # sql counter check
     "tox",  # used for setting up testing environments
+    "snowflake.core>=1.0.0, <2",  # Catalog
     "psutil",  # testing for telemetry
     "lxml",  # used in XML reader unit tests
 ]

src/snowflake/snowpark/_immutable_attr_dict.py

@@ -4,6 +4,7 @@
 import math
 import os
 import queue
+import re
 import time
 import traceback
 import threading
@@ -46,10 +47,15 @@
 _MAX_RETRY_TIME = 3
 _MAX_WORKER_SCALE = 2  # 2 * max_workers
 STATEMENT_PARAMS_DATA_SOURCE = "SNOWPARK_PYTHON_DATASOURCE"
+STATEMENT_PARAMS_DATA_SOURCE_JDBC = "SNOWPARK_PYTHON_DATASOURCE_JDBC"
 DATA_SOURCE_DBAPI_SIGNATURE = "DataFrameReader.dbapi"
+DATA_SOURCE_JDBC_SIGNATURE = "DataFrameReader.jdbc"
 DATA_SOURCE_SQL_COMMENT = (
     f"/* Python:snowflake.snowpark.{DATA_SOURCE_DBAPI_SIGNATURE} */"
 )
+DATA_SOURCE_JDBC_SQL_COMMENT = (
+    f"/* Python:snowflake.snowpark.{DATA_SOURCE_JDBC_SIGNATURE} */"
+)
 PARTITION_TASK_COMPLETE_SIGNAL_PREFIX = "PARTITION_COMPLETE_"
 PARTITION_TASK_ERROR_SIGNAL = "ERROR"
 
@@ -109,6 +115,18 @@ class DRIVER_TYPE(str, Enum):
 }
 
 
+def get_jdbc_dbms(jdbc_url: str) -> str:
+    """
+    Extract the DBMS name from a JDBC connection URL.
+    """
+    if not jdbc_url.startswith("jdbc:"):
+        return "connection url does not start with jdbc"
+
+    # Extract the DBMS type (first component after "jdbc:")
+    match = re.match(r"^jdbc:([^:]+):", jdbc_url)
+    return match.group(1).lower() if match else "unrecognized DBMS"
+
+
 def detect_dbms(dbapi2_conn) -> Tuple[DBMS_TYPE, DRIVER_TYPE]:
     """Detects the DBMS type from a DBAPI2 connection."""