databricks-sql-connector 4.0.3__tar.gz → 4.0.4__tar.gz

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/CHANGELOG.md

@@ -1,5 +1,14 @@
 # Release History
 
+# 4.0.4 (2025-06-16)
+
+- Update thrift client library after cleaning up unused fields and structs (databricks/databricks-sql-python#553 by @vikrantpuppala)
+- Refactor decimal conversion in PyArrow tables to use direct casting (databricks/databricks-sql-python#544 by @jayantsing-db)
+- Fix: `fetchall_arrow` to always return results in `arrow` format (databricks/databricks-sql-python#551 by @shivam2680)
+- Enhance cursor close handling and context manager exception management to prevent server side resource leaks (databricks/databricks-sql-python#554 by @madhav-db)
+- Added additional logging to enhance debugging (databricks/databricks-sql-python#556 by @saishreeeee)
+- Feature: Added support for complex data types such as Arrays and Map [Private Preview] (databricks/databricks-sql-python#559 by @jprakash-db)
+
 # 4.0.3 (2025-04-22)
 
 - Fix: Removed `packaging` dependency in favour of default libraries, for `urllib3` version checks (databricks/databricks-sql-python#547 by @jprakash-db)

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-sql-connector
-Version: 4.0.3
+Version: 4.0.4
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
 Author: Databricks

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.0.3"
+version = "4.0.4"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/src/databricks/sql/__init__.py

@@ -68,7 +68,7 @@ DATETIME = DBAPITypeObject("timestamp")
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "4.0.3"
+__version__ = "4.0.4"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/src/databricks/sql/client.py

@@ -214,6 +214,12 @@ class Connection:
         # use_cloud_fetch
         # Enable use of cloud fetch to extract large query results in parallel via cloud storage
 
+        logger.debug(
+            "Connection.__init__(server_hostname=%s, http_path=%s)",
+            server_hostname,
+            http_path,
+        )
+
         if access_token:
             access_token_kv = {"access_token": access_token}
             kwargs = {**kwargs, **access_token_kv}
@@ -315,7 +321,13 @@ class Connection:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
+        try:
+            self.close()
+        except BaseException as e:
+            logger.warning(f"Exception during connection close in __exit__: {e}")
+            if exc_type is None:
+                raise
+        return False
 
     def __del__(self):
         if self.open:
@@ -456,7 +468,14 @@ class Cursor:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
+        try:
+            logger.debug("Cursor context manager exiting, calling close()")
+            self.close()
+        except BaseException as e:
+            logger.warning(f"Exception during cursor close in __exit__: {e}")
+            if exc_type is None:
+                raise
+        return False
 
     def __iter__(self):
         if self.active_result_set:
@@ -787,6 +806,9 @@ class Cursor:
 
         :returns self
         """
+        logger.debug(
+            "Cursor.execute(operation=%s, parameters=%s)", operation, parameters
+        )
 
         param_approach = self._determine_parameter_approach(parameters)
         if param_approach == ParameterApproach.NONE:
@@ -1163,7 +1185,21 @@ class Cursor:
     def close(self) -> None:
         """Close cursor"""
         self.open = False
-        self.active_op_handle = None
+
+        # Close active operation handle if it exists
+        if self.active_op_handle:
+            try:
+                self.thrift_backend.close_command(self.active_op_handle)
+            except RequestError as e:
+                if isinstance(e.args[1], CursorAlreadyClosedError):
+                    logger.info("Operation was canceled by a prior request")
+                else:
+                    logging.warning(f"Error closing operation handle: {e}")
+            except Exception as e:
+                logging.warning(f"Error closing operation handle: {e}")
+            finally:
+                self.active_op_handle = None
+
         if self.active_result_set:
             self._close_and_clear_active_result_set()
 
@@ -1415,9 +1451,22 @@ class ResultSet:
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
-            results = pyarrow.concat_tables([results, partial_results])
+            if isinstance(results, ColumnTable) and isinstance(
+                partial_results, ColumnTable
+            ):
+                results = self.merge_columnar(results, partial_results)
+            else:
+                results = pyarrow.concat_tables([results, partial_results])
             self._next_row_index += partial_results.num_rows
 
+        # If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table
+        # Valid only for metadata commands result set
+        if isinstance(results, ColumnTable) and pyarrow:
+            data = {
+                name: col
+                for name, col in zip(results.column_names, results.column_table)
+            }
+            return pyarrow.Table.from_pydict(data)
         return results
 
     def fetchall_columnar(self):

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/src/databricks/sql/parameters/__init__.py

@@ -12,4 +12,6 @@ from databricks.sql.parameters.native import (
     TimestampNTZParameter,
     TinyIntParameter,
     DecimalParameter,
+    MapParameter,
+    ArrayParameter,
 )

{databricks_sql_connector-4.0.3 → databricks_sql_connector-4.0.4}/src/databricks/sql/parameters/native.py

@@ -1,12 +1,13 @@
 import datetime
 import decimal
 from enum import Enum, auto
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Any
 
 from databricks.sql.exc import NotSupportedError
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TSparkParameter,
     TSparkParameterValue,
+    TSparkParameterValueArg,
 )
 
 import datetime
@@ -54,7 +55,17 @@ class DatabricksSupportedType(Enum):
 
 
 TAllowedParameterValue = Union[
-    str, int, float, datetime.datetime, datetime.date, bool, decimal.Decimal, None
+    str,
+    int,
+    float,
+    datetime.datetime,
+    datetime.date,
+    bool,
+    decimal.Decimal,
+    None,
+    list,
+    dict,
+    tuple,
 ]
 
 
@@ -82,6 +93,7 @@ class DbsqlParameterBase:
 
     CAST_EXPR: str
     name: Optional[str]
+    value: Any
 
     def as_tspark_param(self, named: bool) -> TSparkParameter:
         """Returns a TSparkParameter object that can be passed to the DBR thrift server."""
@@ -98,6 +110,10 @@ class DbsqlParameterBase:
     def _tspark_param_value(self):
         return TSparkParameterValue(stringValue=str(self.value))
 
+    def _tspark_value_arg(self):
+        """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
+        return TSparkParameterValueArg(value=str(self.value), type=self._cast_expr())
+
     def _cast_expr(self):
         return self.CAST_EXPR
 
@@ -428,6 +444,99 @@ class TinyIntParameter(DbsqlParameterBase):
     CAST_EXPR = DatabricksSupportedType.TINYINT.name
 
 
+class ArrayParameter(DbsqlParameterBase):
+    """Wrap a Python `Sequence` that will be bound to a Databricks SQL ARRAY type."""
+
+    def __init__(self, value: Sequence[Any], name: Optional[str] = None):
+        """
+        :value:
+            The value to bind for this parameter. This will be casted to a ARRAY.
+        :name:
+            If None, your query must contain a `?` marker. Like:
+
+            ```sql
+               SELECT * FROM table WHERE field = ?
+            ```
+            If not None, your query should contain a named parameter marker. Like:
+            ```sql
+                SELECT * FROM table WHERE field = :my_param
+            ```
+
+            The `name` argument to this function would be `my_param`.
+        """
+        self.name = name
+        self.value = [dbsql_parameter_from_primitive(val) for val in value]
+
+    def as_tspark_param(self, named: bool = False) -> TSparkParameter:
+        """Returns a TSparkParameter object that can be passed to the DBR thrift server."""
+
+        tsp = TSparkParameter(type=self._cast_expr())
+        tsp.arguments = [val._tspark_value_arg() for val in self.value]
+
+        if named:
+            tsp.name = self.name
+            tsp.ordinal = False
+        elif not named:
+            tsp.ordinal = True
+        return tsp
+
+    def _tspark_value_arg(self):
+        """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
+        tva = TSparkParameterValueArg(type=self._cast_expr())
+        tva.arguments = [val._tspark_value_arg() for val in self.value]
+        return tva
+
+    CAST_EXPR = DatabricksSupportedType.ARRAY.name
+
+
+class MapParameter(DbsqlParameterBase):
+    """Wrap a Python `dict` that will be bound to a Databricks SQL MAP type."""
+
+    def __init__(self, value: dict, name: Optional[str] = None):
+        """
+        :value:
+            The value to bind for this parameter. This will be casted to a MAP.
+        :name:
+            If None, your query must contain a `?` marker. Like:
+
+            ```sql
+               SELECT * FROM table WHERE field = ?
+            ```
+            If not None, your query should contain a named parameter marker. Like:
+            ```sql
+                SELECT * FROM table WHERE field = :my_param
+            ```
+
+            The `name` argument to this function would be `my_param`.
+        """
+        self.name = name
+        self.value = [
+            dbsql_parameter_from_primitive(item)
+            for key, val in value.items()
+            for item in (key, val)
+        ]
+
+    def as_tspark_param(self, named: bool = False) -> TSparkParameter:
+        """Returns a TSparkParameter object that can be passed to the DBR thrift server."""
+
+        tsp = TSparkParameter(type=self._cast_expr())
+        tsp.arguments = [val._tspark_value_arg() for val in self.value]
+        if named:
+            tsp.name = self.name
+            tsp.ordinal = False
+        elif not named:
+            tsp.ordinal = True
+        return tsp
+
+    def _tspark_value_arg(self):
+        """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
+        tva = TSparkParameterValueArg(type=self._cast_expr())
+        tva.arguments = [val._tspark_value_arg() for val in self.value]
+        return tva
+
+    CAST_EXPR = DatabricksSupportedType.MAP.name
+
+
 class DecimalParameter(DbsqlParameterBase):
     """Wrap a Python `Decimal` that will be bound to a Databricks SQL DECIMAL type."""
 
@@ -543,23 +652,26 @@ def dbsql_parameter_from_primitive(
     # havoc. We can't use TYPE_INFERRENCE_MAP because mypy doesn't trust
     # its logic
 
-    if type(value) is int:
+    if isinstance(value, bool):
+        return BooleanParameter(value=value, name=name)
+    elif isinstance(value, int):
         return dbsql_parameter_from_int(value, name=name)
-    elif type(value) is str:
+    elif isinstance(value, str):
         return StringParameter(value=value, name=name)
-    elif type(value) is float:
+    elif isinstance(value, float):
         return FloatParameter(value=value, name=name)
-    elif type(value) is datetime.datetime:
+    elif isinstance(value, datetime.datetime):
         return TimestampParameter(value=value, name=name)
-    elif type(value) is datetime.date:
+    elif isinstance(value, datetime.date):
         return DateParameter(value=value, name=name)
-    elif type(value) is bool:
-        return BooleanParameter(value=value, name=name)
-    elif type(value) is decimal.Decimal:
+    elif isinstance(value, decimal.Decimal):
         return DecimalParameter(value=value, name=name)
+    elif isinstance(value, dict):
+        return MapParameter(value=value, name=name)
+    elif isinstance(value, Sequence) and not isinstance(value, str):
+        return ArrayParameter(value=value, name=name)
     elif value is None:
         return VoidParameter(value=value, name=name)
-
     else:
         raise NotSupportedError(
             f"Could not infer parameter type from value: {value} - {type(value)} \n"
@@ -581,6 +693,8 @@ TDbsqlParameter = Union[
     TimestampNTZParameter,
     TinyIntParameter,
     DecimalParameter,
+    ArrayParameter,
+    MapParameter,
 ]