snowflake-sqlalchemy 1.7.3__py3-none-any.whl

snowflake/sqlalchemy/custom_types.py

@@ -0,0 +1,155 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from typing import Optional, Tuple, Union
+
+import sqlalchemy.types as sqltypes
+import sqlalchemy.util as util
+from sqlalchemy.types import TypeEngine
+
+TEXT = sqltypes.VARCHAR
+CHARACTER = sqltypes.CHAR
+DEC = sqltypes.DECIMAL
+DOUBLE = sqltypes.FLOAT
+FIXED = sqltypes.DECIMAL
+NUMBER = sqltypes.DECIMAL
+BYTEINT = sqltypes.SMALLINT
+STRING = sqltypes.VARCHAR
+TINYINT = sqltypes.SMALLINT
+VARBINARY = sqltypes.BINARY
+
+
+def _process_float(value):
+    if value == float("inf"):
+        return "inf"
+    elif value == float("-inf"):
+        return "-inf"
+    elif value is not None:
+        return float(value)
+    return value
+
+
+class SnowflakeType(sqltypes.TypeEngine):
+    def _default_dialect(self):
+        # Get around circular import
+        return __import__("snowflake.sqlalchemy").sqlalchemy.dialect()
+
+
+class VARIANT(SnowflakeType):
+    __visit_name__ = "VARIANT"
+
+
+class StructuredType(SnowflakeType):
+    def __init__(self, is_semi_structured: bool = False):
+        self.is_semi_structured = is_semi_structured
+        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(StructuredType):
+    __visit_name__ = "OBJECT"
+
+    def __init__(self, **items_types: Union[TypeEngine, Tuple[TypeEngine, bool]]):
+        for key, value in items_types.items():
+            if not isinstance(value, tuple):
+                items_types[key] = (value, False)
+
+        self.items_types = items_types
+        self.is_semi_structured = len(items_types) == 0
+        super().__init__()
+
+    def __repr__(self):
+        quote_char = "'"
+        return "OBJECT(%s)" % ", ".join(
+            [
+                f"{repr(key).strip(quote_char)}={repr(value)}"
+                for key, value in self.items_types.items()
+            ]
+        )
+
+
+class ARRAY(StructuredType):
+    __visit_name__ = "SNOWFLAKE_ARRAY"
+
+    def __init__(
+        self,
+        value_type: Optional[sqltypes.TypeEngine] = None,
+        not_null: bool = False,
+    ):
+        self.value_type = value_type
+        self.not_null = not_null
+        super().__init__(is_semi_structured=value_type is None)
+
+
+class TIMESTAMP_TZ(SnowflakeType):
+    __visit_name__ = "TIMESTAMP_TZ"
+
+
+class TIMESTAMP_LTZ(SnowflakeType):
+    __visit_name__ = "TIMESTAMP_LTZ"
+
+
+class TIMESTAMP_NTZ(SnowflakeType):
+    __visit_name__ = "TIMESTAMP_NTZ"
+
+
+class GEOGRAPHY(SnowflakeType):
+    __visit_name__ = "GEOGRAPHY"
+
+
+class GEOMETRY(SnowflakeType):
+    __visit_name__ = "GEOMETRY"
+
+
+class _CUSTOM_Date(SnowflakeType, sqltypes.Date):
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                return f"'{value.isoformat()}'"
+
+        return process
+
+
+class _CUSTOM_DateTime(SnowflakeType, sqltypes.DateTime):
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                datetime_str = value.isoformat(" ", timespec="microseconds")
+                return f"'{datetime_str}'"
+
+        return process
+
+
+class _CUSTOM_Time(SnowflakeType, sqltypes.Time):
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                time_str = value.isoformat(timespec="microseconds")
+                return f"'{time_str}'"
+
+        return process
+
+
+class _CUSTOM_Float(SnowflakeType, sqltypes.Float):
+    def bind_processor(self, dialect):
+        return _process_float
+
+
+class _CUSTOM_DECIMAL(SnowflakeType, sqltypes.DECIMAL):
+    @util.memoized_property
+    def _type_affinity(self):
+        return sqltypes.INTEGER if self.scale == 0 else sqltypes.DECIMAL

snowflake/sqlalchemy/exc.py

@@ -0,0 +1,82 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+
+from typing import List
+
+from sqlalchemy.exc import ArgumentError
+
+
+class NoPrimaryKeyError(ArgumentError):
+    def __init__(self, target: str):
+        super().__init__(f"Table {target} required primary key.")
+
+
+class UnsupportedPrimaryKeysAndForeignKeysError(ArgumentError):
+    def __init__(self, target: str):
+        super().__init__(f"Primary key and foreign keys are not supported in {target}.")
+
+
+class RequiredParametersNotProvidedError(ArgumentError):
+    def __init__(self, target: str, parameters: List[str]):
+        super().__init__(
+            f"{target} requires the following parameters: %s." % ", ".join(parameters)
+        )
+
+
+class UnexpectedTableOptionKeyError(ArgumentError):
+    def __init__(self, expected: str, actual: str):
+        super().__init__(f"Expected table option {expected} but got {actual}.")
+
+
+class OptionKeyNotProvidedError(ArgumentError):
+    def __init__(self, target: str):
+        super().__init__(
+            f"Expected option key in {target} option but got NoneType instead."
+        )
+
+
+class UnexpectedOptionParameterTypeError(ArgumentError):
+    def __init__(self, parameter_name: str, target: str, types: List[str]):
+        super().__init__(
+            f"Parameter {parameter_name} of {target} requires to be one"
+            f" of following types: {', '.join(types)}."
+        )
+
+
+class CustomOptionsAreOnlySupportedOnSnowflakeTables(ArgumentError):
+    def __init__(self):
+        super().__init__(
+            "Identifier, Literal, TargetLag and other custom options are only supported on Snowflake tables."
+        )
+
+
+class UnexpectedOptionTypeError(ArgumentError):
+    def __init__(self, options: List[str]):
+        super().__init__(
+            f"The following options are either unsupported or should be defined using a Snowflake table: {', '.join(options)}."
+        )
+
+
+class InvalidTableParameterTypeError(ArgumentError):
+    def __init__(self, name: str, input_type: str, expected_types: List[str]):
+        expected_types_str = "', '".join(expected_types)
+        super().__init__(
+            f"Invalid parameter type '{input_type}' provided for '{name}'. "
+            f"Expected one of the following types: '{expected_types_str}'.\n"
+        )
+
+
+class MultipleErrors(ArgumentError):
+    def __init__(self, errors):
+        self.errors = 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."
+        )

snowflake/sqlalchemy/functions.py

@@ -0,0 +1,16 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+
+import warnings
+
+from sqlalchemy.sql import functions as sqlfunc
+
+FLATTEN_WARNING = "For backward compatibility params are not rendered."
+
+
+class flatten(sqlfunc.GenericFunction):
+    name = "flatten"
+
+    def __init__(self, *args, **kwargs):
+        warnings.warn(FLATTEN_WARNING, DeprecationWarning, stacklevel=2)
+        super().__init__(*args, **kwargs)

snowflake/sqlalchemy/parser/custom_type_parser.py

@@ -0,0 +1,245 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+from typing import List
+
+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 have parameters
+    "INT": INTEGER,
+    "INTEGER": INTEGER,
+    "NUMBER": _CUSTOM_DECIMAL,
+    "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,
+}
+
+NOT_NULL_STR = "NOT NULL"
+
+
+def tokenize_parameters(text: str, character_for_strip=",") -> list:
+    """
+    Extracts parameters from a comma-separated string, handling parentheses.
+
+    :param text: A string with comma-separated parameters, which may include parentheses.
+
+    :param character_for_strip: A character to strip the text.
+
+    :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 != character_for_strip:
+            parameter += c
+        elif c == character_for_strip:
+            output_parameters.append(parameter.strip(" "))
+            parameter = ""
+    if parameter != "":
+        output_parameters.append(parameter.strip(" "))
+    return output_parameters
+
+
+def parse_index_columns(columns: str) -> List[str]:
+    """
+    Parses a string with a list of columns for an index.
+
+    :param columns: A string with a list of columns for an index, which may include parentheses.
+    :param compiler: A SQLAlchemy compiler.
+
+    :return: A list of columns as strings.
+
+    :example:
+        For input `"[A, B, C]"`, the output is `['A', 'B', 'C']`.
+    """
+    return [column.strip() for column in columns.strip("[]").split(",")]
+
+
+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 = (
+        tokenize_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)
+        elif issubclass(col_type_class, OBJECT):
+            col_type_kw = __parse_object_type_parameters(parameters)
+        elif issubclass(col_type_class, ARRAY):
+            col_type_kw = __parse_nullable_parameter(parameters)
+        if col_type_kw is None:
+            col_type_class = NullType
+            col_type_kw = {}
+
+    return col_type_class(**col_type_kw)
+
+
+def __parse_object_type_parameters(parameters):
+    object_rows = {}
+    not_null_parts = NOT_NULL_STR.split(" ")
+    for parameter in parameters:
+        parameter_parts = tokenize_parameters(parameter, " ")
+        if len(parameter_parts) >= 2:
+            key = parameter_parts[0]
+            value_type = parse_type(parameter_parts[1])
+            if isinstance(value_type, NullType):
+                return None
+            not_null = (
+                len(parameter_parts) == 4
+                and parameter_parts[2] == not_null_parts[0]
+                and parameter_parts[3] == not_null_parts[1]
+            )
+            object_rows[key] = (value_type, not_null)
+    return object_rows
+
+
+def __parse_nullable_parameter(parameters):
+    if len(parameters) < 1:
+        return {}
+    elif len(parameters) > 1:
+        return None
+    parameter_str = parameters[0]
+    is_not_null = False
+    if (
+        len(parameter_str) >= len(NOT_NULL_STR)
+        and parameter_str[-len(NOT_NULL_STR) :] == NOT_NULL_STR
+    ):
+        is_not_null = True
+        parameter_str = parameter_str[: -len(NOT_NULL_STR) - 1]
+
+    value_type: TypeEngine = parse_type(parameter_str)
+    if isinstance(value_type, NullType):
+        return None
+
+    return {
+        "value_type": value_type,
+        "not_null": is_not_null,
+    }
+
+
+def __parse_map_type_parameters(parameters):
+    if len(parameters) != 2:
+        return None
+
+    key_type_str = parameters[0]
+    value_type_str = parameters[1]
+    key_type: TypeEngine = parse_type(key_type_str)
+    value_type = __parse_nullable_parameter([value_type_str])
+    if isinstance(value_type, NullType) or isinstance(key_type, NullType):
+        return None
+
+    return {"key_type": key_type, **value_type}
+
+
+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

snowflake/sqlalchemy/provision.py

@@ -0,0 +1,12 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from sqlalchemy.testing.provision import set_default_schema_on_connection
+
+
+# This is only for test purpose required by Requirement "default_schema_name_switch"
+@set_default_schema_on_connection.for_db("snowflake")
+def _snowflake_set_default_schema_on_connection(cfg, dbapi_connection, schema_name):
+    cursor = dbapi_connection.cursor()
+    cursor.execute(f"USE SCHEMA {dbapi_connection.database}.{schema_name};")
+    cursor.close()