snowflake-sqlalchemy 1.7.3__py3-none-any.whl

snowflake/sqlalchemy/sql/__init__.py

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#

snowflake/sqlalchemy/sql/custom_schema/__init__.py

@@ -0,0 +1,9 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from .dynamic_table import DynamicTable
+from .hybrid_table import HybridTable
+from .iceberg_table import IcebergTable
+from .snowflake_table import SnowflakeTable
+
+__all__ = ["DynamicTable", "HybridTable", "IcebergTable", "SnowflakeTable"]

snowflake/sqlalchemy/sql/custom_schema/clustered_table.py

@@ -0,0 +1,37 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Any, Optional
+
+from sqlalchemy.sql.schema import MetaData, SchemaItem
+
+from .custom_table_base import CustomTableBase
+from .options.as_query_option import AsQueryOption
+from .options.cluster_by_option import ClusterByOption, ClusterByOptionType
+from .options.table_option import TableOptionKey
+
+
+class ClusteredTableBase(CustomTableBase):
+
+    @property
+    def cluster_by(self) -> Optional[AsQueryOption]:
+        return self._get_dialect_option(TableOptionKey.CLUSTER_BY)
+
+    def __init__(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        cluster_by: ClusterByOptionType = None,
+        **kw: Any,
+    ) -> None:
+        if kw.get("_no_init", True):
+            return
+
+        options = [
+            ClusterByOption.create(cluster_by),
+        ]
+
+        kw.update(self._as_dialect_options(options))
+        super().__init__(name, metadata, *args, **kw)

snowflake/sqlalchemy/sql/custom_schema/custom_table_base.py

@@ -0,0 +1,127 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+import typing
+from typing import Any, List
+
+from sqlalchemy.sql.schema import MetaData, SchemaItem, Table
+
+from ..._constants import DIALECT_NAME
+from ...compat import IS_VERSION_20
+from ...custom_commands import NoneType
+from ...custom_types import StructuredType
+from ...exc import (
+    MultipleErrors,
+    NoPrimaryKeyError,
+    RequiredParametersNotProvidedError,
+    StructuredTypeNotSupportedInTableColumnsError,
+    UnsupportedPrimaryKeysAndForeignKeysError,
+)
+from .custom_table_prefix import CustomTablePrefix
+from .options.invalid_table_option import InvalidTableOption
+from .options.table_option import TableOption, TableOptionKey
+
+
+class CustomTableBase(Table):
+    __table_prefixes__: typing.List[CustomTablePrefix] = []
+    _support_primary_and_foreign_keys: bool = True
+    _enforce_primary_keys: bool = False
+    _required_parameters: List[TableOptionKey] = []
+    _support_structured_types: bool = False
+
+    @property
+    def table_prefixes(self) -> typing.List[str]:
+        return [prefix.name for prefix in self.__table_prefixes__]
+
+    def __init__(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        **kw: Any,
+    ) -> None:
+        if len(self.__table_prefixes__) > 0:
+            prefixes = kw.get("prefixes", []) + self.table_prefixes
+            kw.update(prefixes=prefixes)
+
+        if not IS_VERSION_20 and hasattr(super(), "_init"):
+            kw.pop("_no_init", True)
+            super()._init(name, metadata, *args, **kw)
+        else:
+            super().__init__(name, metadata, *args, **kw)
+
+        if not kw.get("autoload_with", False):
+            self._validate_table()
+
+    def _validate_table(self):
+        exceptions: List[Exception] = []
+
+        columns_validation = self.__validate_columns()
+        if columns_validation is not None:
+            exceptions.append(columns_validation)
+
+        for _, option in self.dialect_options[DIALECT_NAME].items():
+            if isinstance(option, InvalidTableOption):
+                exceptions.append(option.exception)
+
+        if isinstance(self.key, NoneType) and self._enforce_primary_keys:
+            exceptions.append(NoPrimaryKeyError(self.__class__.__name__))
+        missing_parameters: List[str] = []
+
+        for required_parameter in self._required_parameters:
+            if isinstance(self._get_dialect_option(required_parameter), NoneType):
+                missing_parameters.append(required_parameter.name.lower())
+        if missing_parameters:
+            exceptions.append(
+                RequiredParametersNotProvidedError(
+                    self.__class__.__name__, missing_parameters
+                )
+            )
+
+        if not self._support_primary_and_foreign_keys and (
+            self.primary_key or self.foreign_keys
+        ):
+            exceptions.append(
+                UnsupportedPrimaryKeysAndForeignKeysError(self.__class__.__name__)
+            )
+
+        if len(exceptions) > 1:
+            exceptions.sort(key=lambda e: str(e))
+            raise MultipleErrors(exceptions)
+        elif len(exceptions) == 1:
+            raise exceptions[0]
+
+    def __validate_columns(self):
+        for column in self.columns:
+            if not self._support_structured_types and isinstance(
+                column.type, StructuredType
+            ):
+                return StructuredTypeNotSupportedInTableColumnsError(
+                    self.__class__.__name__, self.name, column.name
+                )
+
+    def _get_dialect_option(
+        self, option_name: TableOptionKey
+    ) -> typing.Optional[TableOption]:
+        if option_name.value in self.dialect_options[DIALECT_NAME]:
+            return self.dialect_options[DIALECT_NAME][option_name.value]
+        return None
+
+    def _as_dialect_options(
+        self, table_options: List[TableOption]
+    ) -> typing.Dict[str, TableOption]:
+        result = {}
+        for table_option in table_options:
+            if isinstance(table_option, TableOption) and isinstance(
+                table_option.option_name, str
+            ):
+                result[DIALECT_NAME + "_" + table_option.option_name] = table_option
+        return result
+
+    @classmethod
+    def is_equal_type(cls, table: Table) -> bool:
+        for prefix in cls.__table_prefixes__:
+            if prefix.name not in table._prefixes:
+                return False
+
+        return True

snowflake/sqlalchemy/sql/custom_schema/custom_table_prefix.py

@@ -0,0 +1,13 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+
+from enum import Enum
+
+
+class CustomTablePrefix(Enum):
+    DEFAULT = 0
+    EXTERNAL = 1
+    EVENT = 2
+    HYBRID = 3
+    ICEBERG = 4
+    DYNAMIC = 5

snowflake/sqlalchemy/sql/custom_schema/dynamic_table.py

@@ -0,0 +1,117 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+import typing
+from typing import Any, Union
+
+from sqlalchemy.sql.schema import MetaData, SchemaItem
+
+from .custom_table_prefix import CustomTablePrefix
+from .options import (
+    IdentifierOption,
+    IdentifierOptionType,
+    KeywordOptionType,
+    TableOptionKey,
+    TargetLagOption,
+    TargetLagOptionType,
+)
+from .options.keyword_option import KeywordOption
+from .table_from_query import TableFromQueryBase
+
+
+class DynamicTable(TableFromQueryBase):
+    """
+    A class representing a dynamic table with configurable options and settings.
+
+    The `DynamicTable` class allows for the creation and querying of tables with
+    specific options, such as `Warehouse` and `TargetLag`.
+
+    While it does not support reflection at this time, it provides a flexible
+    interface for creating dynamic tables and management.
+
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/sql-reference/sql/create-dynamic-table
+
+    Example using option values:
+        DynamicTable(
+        "dynamic_test_table_1",
+        metadata,
+        Column("id", Integer),
+        Column("name", String),
+        target_lag=(1, TimeUnit.HOURS),
+        warehouse='warehouse_name',
+        refresh_mode=SnowflakeKeyword.AUTO
+        as_query="SELECT id, name from test_table_1;"
+    )
+
+    Example using explicit options:
+        DynamicTable(
+        "dynamic_test_table_1",
+        metadata,
+        Column("id", Integer),
+        Column("name", String),
+        target_lag=TargetLag(1, TimeUnit.HOURS),
+        warehouse=Identifier('warehouse_name'),
+        refresh_mode=KeywordOption(SnowflakeKeyword.AUTO)
+        as_query=AsQuery("SELECT id, name from test_table_1;")
+    )
+    """
+
+    __table_prefixes__ = [CustomTablePrefix.DYNAMIC]
+    _support_primary_and_foreign_keys = False
+    _required_parameters = [
+        TableOptionKey.WAREHOUSE,
+        TableOptionKey.AS_QUERY,
+        TableOptionKey.TARGET_LAG,
+    ]
+
+    @property
+    def warehouse(self) -> typing.Optional[IdentifierOption]:
+        return self._get_dialect_option(TableOptionKey.WAREHOUSE)
+
+    @property
+    def target_lag(self) -> typing.Optional[TargetLagOption]:
+        return self._get_dialect_option(TableOptionKey.TARGET_LAG)
+
+    def __init__(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        warehouse: IdentifierOptionType = None,
+        target_lag: Union[TargetLagOptionType, KeywordOptionType] = None,
+        refresh_mode: KeywordOptionType = None,
+        **kw: Any,
+    ) -> None:
+        if kw.get("_no_init", True):
+            return
+
+        options = [
+            IdentifierOption.create(TableOptionKey.WAREHOUSE, warehouse),
+            TargetLagOption.create(target_lag),
+            KeywordOption.create(TableOptionKey.REFRESH_MODE, refresh_mode),
+        ]
+
+        kw.update(self._as_dialect_options(options))
+        super().__init__(name, metadata, *args, **kw)
+
+    def _init(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        **kw: Any,
+    ) -> None:
+        self.__init__(name, metadata, *args, _no_init=False, **kw)
+
+    def __repr__(self) -> str:
+        return "DynamicTable(%s)" % ", ".join(
+            [repr(self.name)]
+            + [repr(self.metadata)]
+            + [repr(x) for x in self.columns]
+            + [repr(self.target_lag)]
+            + [repr(self.warehouse)]
+            + [repr(self.cluster_by)]
+            + [repr(self.as_query)]
+            + [f"{k}={repr(getattr(self, k))}" for k in ["schema"]]
+        )

snowflake/sqlalchemy/sql/custom_schema/hybrid_table.py

@@ -0,0 +1,63 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Any
+
+from sqlalchemy.sql.schema import MetaData, SchemaItem
+
+from .custom_table_base import CustomTableBase
+from .custom_table_prefix import CustomTablePrefix
+
+
+class HybridTable(CustomTableBase):
+    """
+    A class representing a hybrid table with configurable options and settings.
+
+    The `HybridTable` class allows for the creation and querying of OLTP Snowflake Tables .
+
+    While it does not support reflection at this time, it provides a flexible
+    interface for creating hybrid tables and management.
+
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/sql-reference/sql/create-hybrid-table
+
+    Example usage:
+    HybridTable(
+        table_name,
+        metadata,
+        Column("id", Integer, primary_key=True),
+        Column("name", String)
+    )
+    """
+
+    __table_prefixes__ = [CustomTablePrefix.HYBRID]
+    _enforce_primary_keys: bool = True
+    _support_structured_types = True
+
+    def __init__(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        **kw: Any,
+    ) -> None:
+        if kw.get("_no_init", True):
+            return
+        super().__init__(name, metadata, *args, **kw)
+
+    def _init(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        **kw: Any,
+    ) -> None:
+        self.__init__(name, metadata, *args, _no_init=False, **kw)
+
+    def __repr__(self) -> str:
+        return "HybridTable(%s)" % ", ".join(
+            [repr(self.name)]
+            + [repr(self.metadata)]
+            + [repr(x) for x in self.columns]
+            + [f"{k}={repr(getattr(self, k))}" for k in ["schema"]]
+        )

snowflake/sqlalchemy/sql/custom_schema/iceberg_table.py

@@ -0,0 +1,102 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+import typing
+from typing import Any
+
+from sqlalchemy.sql.schema import MetaData, SchemaItem
+
+from .custom_table_prefix import CustomTablePrefix
+from .options import LiteralOption, LiteralOptionType, TableOptionKey
+from .table_from_query import TableFromQueryBase
+
+
+class IcebergTable(TableFromQueryBase):
+    """
+    A class representing an iceberg table with configurable options and settings.
+
+    While it does not support reflection at this time, it provides a flexible
+    interface for creating iceberg tables and management.
+
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/sql-reference/sql/create-iceberg-table
+
+    Example using option values:
+
+        IcebergTable(
+        "dynamic_test_table_1",
+        metadata,
+        Column("id", Integer),
+        Column("name", String),
+        external_volume='my_external_volume',
+        base_location='my_iceberg_table'"
+    )
+
+    Example using explicit options:
+        DynamicTable(
+        "dynamic_test_table_1",
+        metadata,
+        Column("id", Integer),
+        Column("name", String),
+        external_volume=LiteralOption('my_external_volume')
+        base_location=LiteralOption('my_iceberg_table')
+    )
+    """
+
+    __table_prefixes__ = [CustomTablePrefix.ICEBERG]
+    _support_structured_types = True
+
+    @property
+    def external_volume(self) -> typing.Optional[LiteralOption]:
+        return self._get_dialect_option(TableOptionKey.EXTERNAL_VOLUME)
+
+    @property
+    def base_location(self) -> typing.Optional[LiteralOption]:
+        return self._get_dialect_option(TableOptionKey.BASE_LOCATION)
+
+    @property
+    def catalog(self) -> typing.Optional[LiteralOption]:
+        return self._get_dialect_option(TableOptionKey.CATALOG)
+
+    def __init__(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        external_volume: LiteralOptionType = None,
+        base_location: LiteralOptionType = None,
+        **kw: Any,
+    ) -> None:
+        if kw.get("_no_init", True):
+            return
+
+        options = [
+            LiteralOption.create(TableOptionKey.EXTERNAL_VOLUME, external_volume),
+            LiteralOption.create(TableOptionKey.BASE_LOCATION, base_location),
+            LiteralOption.create(TableOptionKey.CATALOG, "SNOWFLAKE"),
+        ]
+
+        kw.update(self._as_dialect_options(options))
+        super().__init__(name, metadata, *args, **kw)
+
+    def _init(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        **kw: Any,
+    ) -> None:
+        self.__init__(name, metadata, *args, _no_init=False, **kw)
+
+    def __repr__(self) -> str:
+        return "IcebergTable(%s)" % ", ".join(
+            [repr(self.name)]
+            + [repr(self.metadata)]
+            + [repr(x) for x in self.columns]
+            + [repr(self.external_volume)]
+            + [repr(self.base_location)]
+            + [repr(self.catalog)]
+            + [repr(self.cluster_by)]
+            + [repr(self.as_query)]
+            + [f"{k}={repr(getattr(self, k))}" for k in ["schema"]]
+        )

snowflake/sqlalchemy/sql/custom_schema/options/__init__.py

@@ -0,0 +1,33 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from .as_query_option import AsQueryOption, AsQueryOptionType
+from .cluster_by_option import ClusterByOption, ClusterByOptionType
+from .identifier_option import IdentifierOption, IdentifierOptionType
+from .keyword_option import KeywordOption, KeywordOptionType
+from .keywords import SnowflakeKeyword
+from .literal_option import LiteralOption, LiteralOptionType
+from .table_option import TableOptionKey
+from .target_lag_option import TargetLagOption, TargetLagOptionType, TimeUnit
+
+__all__ = [
+    # Options
+    "IdentifierOption",
+    "LiteralOption",
+    "KeywordOption",
+    "AsQueryOption",
+    "TargetLagOption",
+    "ClusterByOption",
+    # Enums
+    "TimeUnit",
+    "SnowflakeKeyword",
+    "TableOptionKey",
+    # Types
+    "IdentifierOptionType",
+    "LiteralOptionType",
+    "AsQueryOptionType",
+    "TargetLagOptionType",
+    "KeywordOptionType",
+    "ClusterByOptionType",
+]

snowflake/sqlalchemy/sql/custom_schema/options/as_query_option.py

@@ -0,0 +1,63 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from typing import Optional, Union
+
+from sqlalchemy.sql import Selectable
+
+from snowflake.sqlalchemy.custom_commands import NoneType
+
+from .table_option import Priority, TableOption, TableOptionKey
+
+
+class AsQueryOption(TableOption):
+    """Class to represent an AS clause in tables.
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/sql-reference/sql/create-table#create-table-as-select-also-referred-to-as-ctas
+
+    Example:
+        as_query=AsQueryOption('select name, address from existing_table where name = "test"')
+
+        is equivalent to:
+
+        as select name, address from existing_table where name = "test"
+    """
+
+    def __init__(self, query: Union[str, Selectable]) -> None:
+        super().__init__()
+        self._name: TableOptionKey = TableOptionKey.AS_QUERY
+        self.query = query
+
+    @staticmethod
+    def create(
+        value: Optional[Union["AsQueryOption", str, Selectable]]
+    ) -> "TableOption":
+        if isinstance(value, (NoneType, AsQueryOption)):
+            return value
+        if isinstance(value, (str, Selectable)):
+            return AsQueryOption(value)
+        return TableOption._get_invalid_table_option(
+            TableOptionKey.AS_QUERY,
+            str(type(value).__name__),
+            [AsQueryOption.__name__, str.__name__, Selectable.__name__],
+        )
+
+    def template(self) -> str:
+        return "AS %s"
+
+    @property
+    def priority(self) -> Priority:
+        return Priority.LOWEST
+
+    def __get_expression(self):
+        if isinstance(self.query, Selectable):
+            return self.query.compile(compile_kwargs={"literal_binds": True})
+        return self.query
+
+    def _render(self, compiler) -> str:
+        return self.template() % (self.__get_expression())
+
+    def __repr__(self) -> str:
+        return "AsQueryOption(%s)" % self.__get_expression()
+
+
+AsQueryOptionType = Union[AsQueryOption, str, Selectable]

snowflake/sqlalchemy/sql/custom_schema/options/cluster_by_option.py

@@ -0,0 +1,58 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from typing import List, Union
+
+from sqlalchemy.sql.expression import TextClause
+
+from snowflake.sqlalchemy.custom_commands import NoneType
+
+from .table_option import Priority, TableOption, TableOptionKey
+
+
+class ClusterByOption(TableOption):
+    """Class to represent the cluster by clause in tables.
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/user-guide/tables-clustering-keys
+    Example:
+        cluster_by=ClusterByOption('name', text('id > 0'))
+
+        is equivalent to:
+
+        cluster by (name, id > 0)
+    """
+
+    def __init__(self, *expressions: Union[str, TextClause]) -> None:
+        super().__init__()
+        self._name: TableOptionKey = TableOptionKey.CLUSTER_BY
+        self.expressions = expressions
+
+    @staticmethod
+    def create(value: "ClusterByOptionType") -> "TableOption":
+        if isinstance(value, (NoneType, ClusterByOption)):
+            return value
+        if isinstance(value, List):
+            return ClusterByOption(*value)
+        return TableOption._get_invalid_table_option(
+            TableOptionKey.CLUSTER_BY,
+            str(type(value).__name__),
+            [ClusterByOption.__name__, list.__name__],
+        )
+
+    def template(self) -> str:
+        return f"{self.option_name.upper()} (%s)"
+
+    @property
+    def priority(self) -> Priority:
+        return Priority.HIGH
+
+    def __get_expression(self):
+        return ", ".join([str(expression) for expression in self.expressions])
+
+    def _render(self, compiler) -> str:
+        return self.template() % (self.__get_expression())
+
+    def __repr__(self) -> str:
+        return "ClusterByOption(%s)" % self.__get_expression()
+
+
+ClusterByOptionType = Union[ClusterByOption, List[Union[str, TextClause]]]

snowflake/sqlalchemy/sql/custom_schema/options/identifier_option.py

@@ -0,0 +1,63 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from typing import Optional, Union
+
+from snowflake.sqlalchemy.custom_commands import NoneType
+
+from .table_option import Priority, TableOption, TableOptionKey
+
+
+class IdentifierOption(TableOption):
+    """Class to represent an identifier option in Snowflake Tables.
+
+    Example:
+        warehouse = IdentifierOption('my_warehouse')
+
+        is equivalent to:
+
+        WAREHOUSE = my_warehouse
+    """
+
+    def __init__(self, value: Union[str]) -> None:
+        super().__init__()
+        self.value: str = value
+
+    @property
+    def priority(self):
+        return Priority.HIGH
+
+    @staticmethod
+    def create(
+        name: TableOptionKey, value: Optional[Union[str, "IdentifierOption"]]
+    ) -> Optional[TableOption]:
+        if isinstance(value, NoneType):
+            return None
+
+        if isinstance(value, str):
+            value = IdentifierOption(value)
+
+        if isinstance(value, IdentifierOption):
+            value._set_option_name(name)
+            return value
+
+        return TableOption._get_invalid_table_option(
+            name, str(type(value).__name__), [IdentifierOption.__name__, str.__name__]
+        )
+
+    def template(self) -> str:
+        return f"{self.option_name.upper()} = %s"
+
+    def _render(self, compiler) -> str:
+        return self.template() % self.value
+
+    def __repr__(self) -> str:
+        option_name = (
+            f", table_option_key={self.option_name}"
+            if not isinstance(self.option_name, NoneType)
+            else ""
+        )
+        return f"IdentifierOption(value='{self.value}'{option_name})"
+
+
+IdentifierOptionType = Union[IdentifierOption, str, int]