Add support for map datatype

DESCRIPTION.md

@@ -9,7 +9,7 @@ Source code is also available at:
 
 # Release Notes
 
-- (Unreleased)
+- v1.7.0(November 12, 2024)
 
   - Add support for dynamic tables and required options
   - Add support for hybrid tables
@@ -18,6 +18,7 @@ Source code is also available at:
   - Add support for refresh_mode option in DynamicTable
   - Add support for iceberg table with Snowflake Catalog
   - Fix cluster by option to support explicit expressions
+  - Add support for MAP datatype
 
 - v1.6.1(July 9, 2024)
 

pyproject.toml

@@ -111,7 +111,7 @@ line-length = 88
 line-length = 88
 
 [tool.pytest.ini_options]
-addopts = "-m 'not feature_max_lob_size and not aws'"
+addopts = "-m 'not feature_max_lob_size and not aws and not requires_external_volume'"
 markers = [
   # Optional dependency groups markers
   "lambda: AWS lambda tests",
@@ -128,6 +128,7 @@ markers = [
   # Other markers
   "timeout: tests that need a timeout time",
   "internal: tests that could but should only run on our internal CI",
+  "requires_external_volume: tests that needs a external volume to be executed",
   "external: tests that could but should only run on our external CI",
   "feature_max_lob_size: tests that could but should only run on our external CI",
 ]

src/snowflake/sqlalchemy/__init__.py

@@ -50,6 +50,7 @@
     FIXED,
     GEOGRAPHY,
     GEOMETRY,
+    MAP,
     NUMBER,
     OBJECT,
     STRING,
@@ -119,6 +120,7 @@
     "TINYINT",
     "VARBINARY",
     "VARIANT",
+    "MAP",
 )
 
 _custom_commands = (

src/snowflake/sqlalchemy/_constants.py

@@ -11,3 +11,4 @@
 APPLICATION_NAME = "SnowflakeSQLAlchemy"
 SNOWFLAKE_SQLALCHEMY_VERSION = VERSION
 DIALECT_NAME = "snowflake"
+NOT_NULL = "NOT NULL"

src/snowflake/sqlalchemy/base.py

@@ -27,6 +27,7 @@
     ExternalStage,
 )
 
+from ._constants import NOT_NULL
 from .exc import (
     CustomOptionsAreOnlySupportedOnSnowflakeTables,
     UnexpectedOptionTypeError,
@@ -1071,6 +1072,12 @@ def visit_TINYINT(self, type_, **kw):
     def visit_VARIANT(self, type_, **kw):
         return "VARIANT"
 
+    def visit_MAP(self, type_, **kw):
+        not_null = f" {NOT_NULL}" if type_.not_null else ""
+        return (
+            f"MAP({type_.key_type.compile()}, {type_.value_type.compile()}{not_null})"
+        )
+
     def visit_ARRAY(self, type_, **kw):
         return "ARRAY"
 

src/snowflake/sqlalchemy/custom_types.py

@@ -37,6 +37,26 @@ class VARIANT(SnowflakeType):
     __visit_name__ = "VARIANT"
 
 
+class StructuredType(SnowflakeType):
+    def __init__(self):
+        super().__init__()
+
+
+class MAP(StructuredType):
+    __visit_name__ = "MAP"
+
+    def __init__(
+        self,
+        key_type: sqltypes.TypeEngine,
+        value_type: sqltypes.TypeEngine,
+        not_null: bool = False,
+    ):
+        self.key_type = key_type
+        self.value_type = value_type
+        self.not_null = not_null
+        super().__init__()
+
+
 class OBJECT(SnowflakeType):
     __visit_name__ = "OBJECT"
 

src/snowflake/sqlalchemy/exc.py

@@ -72,3 +72,11 @@ def __init__(self, errors):
 
     def __str__(self):
         return "".join(str(e) for e in self.errors)
+
+
+class StructuredTypeNotSupportedInTableColumnsError(ArgumentError):
+    def __init__(self, table_type: str, table_name: str, column_name: str):
+        super().__init__(
+            f"Column '{column_name}' is of a structured type, which is only supported on Iceberg tables. "
+            f"The table '{table_name}' is of type '{table_type}', not Iceberg."
+        )

src/snowflake/sqlalchemy/parser/custom_type_parser.py

@@ -0,0 +1,190 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+
+import sqlalchemy.types as sqltypes
+from sqlalchemy.sql.type_api import TypeEngine
+from sqlalchemy.types import (
+    BIGINT,
+    BINARY,
+    BOOLEAN,
+    CHAR,
+    DATE,
+    DATETIME,
+    DECIMAL,
+    FLOAT,
+    INTEGER,
+    REAL,
+    SMALLINT,
+    TIME,
+    TIMESTAMP,
+    VARCHAR,
+    NullType,
+)
+
+from ..custom_types import (
+    _CUSTOM_DECIMAL,
+    ARRAY,
+    DOUBLE,
+    GEOGRAPHY,
+    GEOMETRY,
+    MAP,
+    OBJECT,
+    TIMESTAMP_LTZ,
+    TIMESTAMP_NTZ,
+    TIMESTAMP_TZ,
+    VARIANT,
+)
+
+ischema_names = {
+    "BIGINT": BIGINT,
+    "BINARY": BINARY,
+    # 'BIT': BIT,
+    "BOOLEAN": BOOLEAN,
+    "CHAR": CHAR,
+    "CHARACTER": CHAR,
+    "DATE": DATE,
+    "DATETIME": DATETIME,
+    "DEC": DECIMAL,
+    "DECIMAL": DECIMAL,
+    "DOUBLE": DOUBLE,
+    "FIXED": DECIMAL,
+    "FLOAT": FLOAT,  # Snowflake FLOAT datatype doesn't has parameters
+    "INT": INTEGER,
+    "INTEGER": INTEGER,
+    "NUMBER": _CUSTOM_DECIMAL,
+    # 'OBJECT': ?
+    "REAL": REAL,
+    "BYTEINT": SMALLINT,
+    "SMALLINT": SMALLINT,
+    "STRING": VARCHAR,
+    "TEXT": VARCHAR,
+    "TIME": TIME,
+    "TIMESTAMP": TIMESTAMP,
+    "TIMESTAMP_TZ": TIMESTAMP_TZ,
+    "TIMESTAMP_LTZ": TIMESTAMP_LTZ,
+    "TIMESTAMP_NTZ": TIMESTAMP_NTZ,
+    "TINYINT": SMALLINT,
+    "VARBINARY": BINARY,
+    "VARCHAR": VARCHAR,
+    "VARIANT": VARIANT,
+    "MAP": MAP,
+    "OBJECT": OBJECT,
+    "ARRAY": ARRAY,
+    "GEOGRAPHY": GEOGRAPHY,
+    "GEOMETRY": GEOMETRY,
+}
+
+
+def extract_parameters(text: str) -> list:
+    """
+    Extracts parameters from a comma-separated string, handling parentheses.
+
+    :param text: A string with comma-separated parameters, which may include parentheses.
+
+    :return: A list of parameters as strings.
+
+    :example:
+        For input `"a, (b, c), d"`, the output is `['a', '(b, c)', 'd']`.
+    """
+
+    output_parameters = []
+    parameter = ""
+    open_parenthesis = 0
+    for c in text:
+
+        if c == "(":
+            open_parenthesis += 1
+        elif c == ")":
+            open_parenthesis -= 1
+
+        if open_parenthesis > 0 or c != ",":
+            parameter += c
+        elif c == ",":
+            output_parameters.append(parameter.strip(" "))
+            parameter = ""
+    if parameter != "":
+        output_parameters.append(parameter.strip(" "))
+    return output_parameters
+
+
+def parse_type(type_text: str) -> TypeEngine:
+    """
+    Parses a type definition string and returns the corresponding SQLAlchemy type.
+
+    The function handles types with or without parameters, such as `VARCHAR(255)` or `INTEGER`.
+
+    :param type_text: A string representing a SQLAlchemy type, which may include parameters
+                       in parentheses (e.g., "VARCHAR(255)" or "DECIMAL(10, 2)").
+    :return: An instance of the corresponding SQLAlchemy type class (e.g., `String`, `Integer`),
+             or `NullType` if the type is not recognized.
+
+    :example:
+        parse_type("VARCHAR(255)")
+        String(length=255)
+    """
+    index = type_text.find("(")
+    type_name = type_text[:index] if index != -1 else type_text
+    parameters = (
+        extract_parameters(type_text[index + 1 : -1]) if type_name != type_text else []
+    )
+
+    col_type_class = ischema_names.get(type_name, None)
+    col_type_kw = {}
+    if col_type_class is None:
+        col_type_class = NullType
+    else:
+        if issubclass(col_type_class, sqltypes.Numeric):
+            col_type_kw = __parse_numeric_type_parameters(parameters)
+        elif issubclass(col_type_class, (sqltypes.String, sqltypes.BINARY)):
+            col_type_kw = __parse_type_with_length_parameters(parameters)
+        elif issubclass(col_type_class, MAP):
+            col_type_kw = __parse_map_type_parameters(parameters)
+        if col_type_kw is None:
+            col_type_class = NullType
+            col_type_kw = {}
+
+    return col_type_class(**col_type_kw)
+
+
+def __parse_map_type_parameters(parameters):
+    if len(parameters) != 2:
+        return None
+
+    key_type_str = parameters[0]
+    value_type_str = parameters[1]
+    not_null_str = "NOT NULL"
+    not_null = False
+    if (
+        len(value_type_str) >= len(not_null_str)
+        and value_type_str[-len(not_null_str) :] == not_null_str
+    ):
+        not_null = True
+        value_type_str = value_type_str[: -len(not_null_str) - 1]
+
+    key_type: TypeEngine = parse_type(key_type_str)
+    value_type: TypeEngine = parse_type(value_type_str)
+    if isinstance(key_type, NullType) or isinstance(value_type, NullType):
+        return None
+
+    return {
+        "key_type": key_type,
+        "value_type": value_type,
+        "not_null": not_null,
+    }
+
+
+def __parse_type_with_length_parameters(parameters):
+    return (
+        {"length": int(parameters[0])}
+        if len(parameters) == 1 and str.isdigit(parameters[0])
+        else {}
+    )
+
+
+def __parse_numeric_type_parameters(parameters):
+    result = {}
+    if len(parameters) >= 1 and str.isdigit(parameters[0]):
+        result["precision"] = int(parameters[0])
+    if len(parameters) == 2 and str.isdigit(parameters[1]):
+        result["scale"] = int(parameters[1])
+    return result