databricks-sqlalchemy 2.0.6__tar.gz → 2.0.8__tar.gz

{databricks_sqlalchemy-2.0.6 → databricks_sqlalchemy-2.0.8}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Release History
 
+# 2.0.8 (2025-09-08)
+
+- Feature: Added support for variant datatype (databricks/databricks-sqlalchemy#42 by @msrathore-db)
+
+# 2.0.7 (2025-06-23)
+
+- Feature: Added support for complex data types such as DatabricksArray and DatabricksMap [Private Preview] (databricks/databricks-sqlalchemy#30 by @jprakash-db)
+
 # 2.0.6 (2025-04-29)
 
 - Relaxed pin for `pyarrow` (databricks/databricks-sqlalchemy#20 by @dhirschfeld)

{databricks_sqlalchemy-2.0.6 → databricks_sqlalchemy-2.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: databricks-sqlalchemy
-Version: 2.0.6
+Version: 2.0.8
 Summary: Databricks SQLAlchemy plugin for Python
 License: Apache-2.0
 Author: Databricks
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: databricks_sql_connector (>=4.0.0)
 Requires-Dist: pyarrow (>=14.0.1)
 Requires-Dist: sqlalchemy (>=2.0.21)

{databricks_sqlalchemy-2.0.6 → databricks_sqlalchemy-2.0.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sqlalchemy"
-version = "2.0.6"
+version = "2.0.8"
 description = "Databricks SQLAlchemy plugin for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

databricks_sqlalchemy-2.0.8/src/databricks/sqlalchemy/__init__.py

@@ -0,0 +1,18 @@
+from databricks.sqlalchemy.base import DatabricksDialect
+from databricks.sqlalchemy._types import (
+    TINYINT,
+    TIMESTAMP,
+    TIMESTAMP_NTZ,
+    DatabricksArray,
+    DatabricksMap,
+    DatabricksVariant,
+)
+
+__all__ = [
+    "TINYINT",
+    "TIMESTAMP",
+    "TIMESTAMP_NTZ",
+    "DatabricksArray",
+    "DatabricksMap",
+    "DatabricksVariant",
+]

{databricks_sqlalchemy-2.0.6 → databricks_sqlalchemy-2.0.8}/src/databricks/sqlalchemy/_parse.py

@@ -318,6 +318,7 @@ GET_COLUMNS_TYPE_MAP = {
     "map": sqlalchemy.types.String,
     "struct": sqlalchemy.types.String,
     "uniontype": sqlalchemy.types.String,
+    "variant": type_overrides.DatabricksVariant,
     "decimal": sqlalchemy.types.Numeric,
     "timestamp": type_overrides.TIMESTAMP,
     "timestamp_ntz": type_overrides.TIMESTAMP_NTZ,

{databricks_sqlalchemy-2.0.6 → databricks_sqlalchemy-2.0.8}/src/databricks/sqlalchemy/_types.py

@@ -5,9 +5,13 @@ from typing import Any, Union, Optional
 import sqlalchemy
 from sqlalchemy.engine.interfaces import Dialect
 from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.types import TypeDecorator, UserDefinedType
 
 from databricks.sql.utils import ParamEscaper
 
+from sqlalchemy.sql import expression
+import json
+
 
 def process_literal_param_hack(value: Any):
     """This method is supposed to accept a Python type and return a string representation of that type.
@@ -26,6 +30,11 @@ def process_literal_param_hack(value: Any):
     return value
 
 
+def identity_processor(value):
+    """This method returns the value itself, when no other processor is provided"""
+    return value
+
+
 @compiles(sqlalchemy.types.Enum, "databricks")
 @compiles(sqlalchemy.types.String, "databricks")
 @compiles(sqlalchemy.types.Text, "databricks")
@@ -321,3 +330,130 @@ class TINYINT(sqlalchemy.types.TypeDecorator):
 @compiles(TINYINT, "databricks")
 def compile_tinyint(type_, compiler, **kw):
     return "TINYINT"
+
+
+class DatabricksArray(UserDefinedType):
+    """
+    A custom array type that can wrap any other SQLAlchemy type.
+
+    Examples:
+        DatabricksArray(String)         -> ARRAY<STRING>
+        DatabricksArray(Integer)        -> ARRAY<INT>
+        DatabricksArray(CustomType)     -> ARRAY<CUSTOM_TYPE>
+    """
+
+    def __init__(self, item_type):
+        self.item_type = item_type() if isinstance(item_type, type) else item_type
+
+    def bind_processor(self, dialect):
+        item_processor = self.item_type.bind_processor(dialect)
+        if item_processor is None:
+            item_processor = identity_processor
+
+        def process(value):
+            return [item_processor(val) for val in value]
+
+        return process
+
+
+@compiles(DatabricksArray, "databricks")
+def compile_databricks_array(type_, compiler, **kw):
+    inner = compiler.process(type_.item_type, **kw)
+
+    return f"ARRAY<{inner}>"
+
+
+class DatabricksMap(UserDefinedType):
+    """
+    A custom map type that can wrap any other SQLAlchemy types for both key and value.
+
+    Examples:
+        DatabricksMap(String, String)         -> MAP<STRING,STRING>
+        DatabricksMap(Integer, String)        -> MAP<INT,STRING>
+        DatabricksMap(String, DatabricksArray(Integer)) -> MAP<STRING,ARRAY<INT>>
+    """
+
+    def __init__(self, key_type, value_type):
+        self.key_type = key_type() if isinstance(key_type, type) else key_type
+        self.value_type = value_type() if isinstance(value_type, type) else value_type
+
+    def bind_processor(self, dialect):
+        key_processor = self.key_type.bind_processor(dialect)
+        value_processor = self.value_type.bind_processor(dialect)
+
+        if key_processor is None:
+            key_processor = identity_processor
+        if value_processor is None:
+            value_processor = identity_processor
+
+        def process(value):
+            return {
+                key_processor(key): value_processor(value)
+                for key, value in value.items()
+            }
+
+        return process
+
+
+@compiles(DatabricksMap, "databricks")
+def compile_databricks_map(type_, compiler, **kw):
+    key_type = compiler.process(type_.key_type, **kw)
+    value_type = compiler.process(type_.value_type, **kw)
+    return f"MAP<{key_type},{value_type}>"
+
+
+class DatabricksVariant(UserDefinedType):
+    """
+    A custom variant type for storing semi-structured data including STRUCT, ARRAY, MAP, and scalar types.
+    Note: VARIANT MAP types can only have STRING keys.
+
+    Examples:
+        DatabricksVariant()  -> VARIANT
+
+    Usage:
+        Column('data', DatabricksVariant())
+    """
+
+    cache_ok = True
+
+    def __init__(self):
+        self.pe = ParamEscaper()
+
+    def bind_processor(self, dialect):
+        """Process values before sending to database."""
+
+        def process(value):
+            if value is None:
+                return None
+            try:
+                return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
+            except (TypeError, ValueError) as e:
+                raise ValueError(f"Cannot serialize value {value} to JSON: {e}")
+
+        return process
+
+    def bind_expression(self, bindvalue):
+        """Wrap with PARSE_JSON() in SQL"""
+        return expression.func.PARSE_JSON(bindvalue)
+
+    def literal_processor(self, dialect):
+        """Process literal values for SQL generation.
+        For VARIANT columns, use PARSE_JSON() to properly insert data.
+        """
+
+        def process(value):
+            if value is None:
+                return "NULL"
+            try:
+                return self.pe.escape_string(
+                    json.dumps(value, ensure_ascii=False, separators=(",", ":"))
+                )
+            except (TypeError, ValueError) as e:
+                raise ValueError(f"Cannot serialize value {value} to JSON: {e}")
+
+        return process
+
+
+@compiles(DatabricksVariant, "databricks")
+def compile_variant(type_, compiler, **kw):
+    return "VARIANT"

databricks_sqlalchemy-2.0.6/src/databricks/sqlalchemy/__init__.py

@@ -1,4 +0,0 @@
-from databricks.sqlalchemy.base import DatabricksDialect
-from databricks.sqlalchemy._types import TINYINT, TIMESTAMP, TIMESTAMP_NTZ
-
-__all__ = ["TINYINT", "TIMESTAMP", "TIMESTAMP_NTZ"]