databricks-sqlalchemy 2.0.7__py3-none-any.whl → 2.0.8__py3-none-any.whl

databricks/sqlalchemy/__init__.py

@@ -5,6 +5,14 @@ from databricks.sqlalchemy._types import (
     TIMESTAMP_NTZ,
     DatabricksArray,
     DatabricksMap,
+    DatabricksVariant,
 )
 
-__all__ = ["TINYINT", "TIMESTAMP", "TIMESTAMP_NTZ", "DatabricksArray", "DatabricksMap"]
+__all__ = [
+    "TINYINT",
+    "TIMESTAMP",
+    "TIMESTAMP_NTZ",
+    "DatabricksArray",
+    "DatabricksMap",
+    "DatabricksVariant",
+]

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/_types.py

@@ -9,6 +9,9 @@ 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.
@@ -397,3 +400,60 @@ 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.7.dist-info → databricks_sqlalchemy-2.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databricks-sqlalchemy
-Version: 2.0.7
+Version: 2.0.8
 Summary: Databricks SQLAlchemy plugin for Python
 License: Apache-2.0
 Author: Databricks

{databricks_sqlalchemy-2.0.7.dist-info → databricks_sqlalchemy-2.0.8.dist-info}/RECORD

@@ -1,15 +1,15 @@
 databricks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-databricks/sqlalchemy/__init__.py,sha256=IL3_QTbL9vbdQQ09Dzrol352VXU8v3oUilpiCkjvxQw,278
+databricks/sqlalchemy/__init__.py,sha256=PArcuIQ1nGvovXlY3P-1ly766kh_9MgXsXB_Qgoglp4,349
 databricks/sqlalchemy/_ddl.py,sha256=c0_GwfmnrFVr4-Ls14fmdGUUFyUok_GW4Uo45hLABFc,3983
-databricks/sqlalchemy/_parse.py,sha256=aFpwcLowSDP1R7BY3G-yuEXiPFL-_VaIGvqKNDMehcQ,13049
-databricks/sqlalchemy/_types.py,sha256=nESbFFzf0qFYyomOOBnCAvRB710ABkPgxWMu7qc3mvw,14216
+databricks/sqlalchemy/_parse.py,sha256=oGhErqLLDyGNvVOybfIqYAMJoAa9FhjU6hQODAD8a1w,13098
+databricks/sqlalchemy/_types.py,sha256=SRK0OFn4OORilxZMAnA6QlQ54OGJnt81THTZvRAGa6I,15953
 databricks/sqlalchemy/base.py,sha256=KcjfHMH0NsceYE2NRxrePtf5T1uw9u8JHofRdbnAKS4,15619
 databricks/sqlalchemy/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databricks/sqlalchemy/pytest.ini,sha256=ImutflUjkhByVNWCQ18Todj6XTvgJAQX_v7fD-gWHhU,106
 databricks/sqlalchemy/requirements.py,sha256=OobunAEwZ9y2dvSQLOmdgJciVn9xGlY9NAFfszPCTU0,9018
 databricks/sqlalchemy/setup.cfg,sha256=ImutflUjkhByVNWCQ18Todj6XTvgJAQX_v7fD-gWHhU,106
-databricks_sqlalchemy-2.0.7.dist-info/LICENSE,sha256=WgVm2VpfZ3CsUfPndD2NeCrEIcFA4UB-YnnW4ejxcbE,11346
-databricks_sqlalchemy-2.0.7.dist-info/METADATA,sha256=QIT9EVR2yGQy5N5Z8uDJ9TaMw0li8OXZ1PAS-U2WB5I,12764
-databricks_sqlalchemy-2.0.7.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-databricks_sqlalchemy-2.0.7.dist-info/entry_points.txt,sha256=AAjpsvZbVcoMAcWLIesoAT5FNZhBEcIhxdKknVua3jw,74
-databricks_sqlalchemy-2.0.7.dist-info/RECORD,,
+databricks_sqlalchemy-2.0.8.dist-info/LICENSE,sha256=WgVm2VpfZ3CsUfPndD2NeCrEIcFA4UB-YnnW4ejxcbE,11346
+databricks_sqlalchemy-2.0.8.dist-info/METADATA,sha256=YgQ-w-A6zMEVSeTneL_YHJUK1VmUuj7-c4CCUZqBh04,12764
+databricks_sqlalchemy-2.0.8.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+databricks_sqlalchemy-2.0.8.dist-info/entry_points.txt,sha256=AAjpsvZbVcoMAcWLIesoAT5FNZhBEcIhxdKknVua3jw,74
+databricks_sqlalchemy-2.0.8.dist-info/RECORD,,