snowflake-sqlalchemy 1.7.3__py3-none-any.whl

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

@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from typing import Optional
+
+from .table_option import TableOption, TableOptionKey
+
+
+class InvalidTableOption(TableOption):
+    """Class to store errors and raise them after table initialization in order to avoid recursion error."""
+
+    def __init__(self, name: TableOptionKey, value: Exception) -> None:
+        super().__init__()
+        self.exception: Exception = value
+        self._name = name
+
+    @staticmethod
+    def create(name: TableOptionKey, value: Exception) -> Optional[TableOption]:
+        return InvalidTableOption(name, value)
+
+    def _render(self, compiler) -> str:
+        raise self.exception
+
+    def __repr__(self) -> str:
+        return f"ErrorOption(value='{self.exception}')"

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

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

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

@@ -0,0 +1,14 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+
+from enum import Enum
+
+
+class SnowflakeKeyword(Enum):
+    # TARGET_LAG
+    DOWNSTREAM = "DOWNSTREAM"
+
+    # REFRESH_MODE
+    AUTO = "AUTO"
+    FULL = "FULL"
+    INCREMENTAL = "INCREMENTAL"

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

@@ -0,0 +1,67 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from typing import Any, Optional, Union
+
+from snowflake.sqlalchemy.custom_commands import NoneType
+
+from .table_option import Priority, TableOption, TableOptionKey
+
+
+class LiteralOption(TableOption):
+    """Class to represent a literal option in Snowflake Table.
+
+    Example:
+        warehouse = LiteralOption('my_warehouse')
+
+        is equivalent to:
+
+        WAREHOUSE = 'my_warehouse'
+    """
+
+    def __init__(self, value: Union[int, str]) -> None:
+        super().__init__()
+        self.value: Any = value
+
+    @property
+    def priority(self):
+        return Priority.HIGH
+
+    @staticmethod
+    def create(
+        name: TableOptionKey, value: Optional[Union[str, int, "LiteralOption"]]
+    ) -> Optional[TableOption]:
+        if isinstance(value, NoneType):
+            return None
+        if isinstance(value, (str, int)):
+            value = LiteralOption(value)
+
+        if isinstance(value, LiteralOption):
+            value._set_option_name(name)
+            return value
+
+        return TableOption._get_invalid_table_option(
+            name,
+            str(type(value).__name__),
+            [LiteralOption.__name__, str.__name__, int.__name__],
+        )
+
+    def template(self) -> str:
+        if isinstance(self.value, int):
+            return f"{self.option_name.upper()} = %d"
+        else:
+            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"LiteralOption(value='{self.value}'{option_name})"
+
+
+LiteralOptionType = Union[LiteralOption, str, int]

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

@@ -0,0 +1,84 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+from enum import Enum
+from typing import List, Optional
+
+from snowflake.sqlalchemy import exc
+from snowflake.sqlalchemy.custom_commands import NoneType
+
+
+class Priority(Enum):
+    LOWEST = 0
+    VERY_LOW = 1
+    LOW = 2
+    MEDIUM = 4
+    HIGH = 6
+    VERY_HIGH = 7
+    HIGHEST = 8
+
+
+class TableOption:
+
+    def __init__(self) -> None:
+        self._name: Optional[TableOptionKey] = None
+
+    @property
+    def option_name(self) -> str:
+        if isinstance(self._name, NoneType):
+            return None
+        return str(self._name.value)
+
+    def _set_option_name(self, name: Optional["TableOptionKey"]):
+        self._name = name
+
+    @property
+    def priority(self) -> Priority:
+        return Priority.MEDIUM
+
+    @staticmethod
+    def create(**kwargs) -> "TableOption":
+        raise NotImplementedError
+
+    @staticmethod
+    def _get_invalid_table_option(
+        parameter_name: "TableOptionKey", input_type: str, expected_types: List[str]
+    ) -> "TableOption":
+        from .invalid_table_option import InvalidTableOption
+
+        return InvalidTableOption(
+            parameter_name,
+            exc.InvalidTableParameterTypeError(
+                parameter_name.value, input_type, expected_types
+            ),
+        )
+
+    def _validate_option(self):
+        if isinstance(self.option_name, NoneType):
+            raise exc.OptionKeyNotProvidedError(self.__class__.__name__)
+
+    def template(self) -> str:
+        return f"{self.option_name.upper()} = %s"
+
+    def render_option(self, compiler) -> str:
+        self._validate_option()
+        return self._render(compiler)
+
+    def _render(self, compiler) -> str:
+        raise NotImplementedError
+
+
+class TableOptionKey(Enum):
+    AS_QUERY = "as_query"
+    BASE_LOCATION = "base_location"
+    CATALOG = "catalog"
+    CATALOG_SYNC = "catalog_sync"
+    CLUSTER_BY = "cluster by"
+    DATA_RETENTION_TIME_IN_DAYS = "data_retention_time_in_days"
+    DEFAULT_DDL_COLLATION = "default_ddl_collation"
+    EXTERNAL_VOLUME = "external_volume"
+    MAX_DATA_EXTENSION_TIME_IN_DAYS = "max_data_extension_time_in_days"
+    REFRESH_MODE = "refresh_mode"
+    STORAGE_SERIALIZATION_POLICY = "storage_serialization_policy"
+    TARGET_LAG = "target_lag"
+    WAREHOUSE = "warehouse"

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

@@ -0,0 +1,94 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+# from enum import Enum
+from enum import Enum
+from typing import Optional, Tuple, Union
+
+from snowflake.sqlalchemy.custom_commands import NoneType
+
+from .keyword_option import KeywordOption, KeywordOptionType
+from .keywords import SnowflakeKeyword
+from .table_option import Priority, TableOption, TableOptionKey
+
+
+class TimeUnit(Enum):
+    SECONDS = "seconds"
+    MINUTES = "minutes"
+    HOURS = "hours"
+    DAYS = "days"
+
+
+class TargetLagOption(TableOption):
+    """Class to represent the target lag clause in Dynamic Tables.
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/sql-reference/sql/create-dynamic-table
+
+    Example using the time and unit parameters:
+
+        target_lag = TargetLagOption(10, TimeUnit.SECONDS)
+
+        is equivalent to:
+
+        TARGET_LAG = '10 SECONDS'
+
+    Example using keyword parameter:
+
+        target_lag = KeywordOption(SnowflakeKeyword.DOWNSTREAM)
+
+        is equivalent to:
+
+        TARGET_LAG = DOWNSTREAM
+
+    """
+
+    def __init__(
+        self,
+        time: Optional[int] = 0,
+        unit: Optional[TimeUnit] = TimeUnit.MINUTES,
+    ) -> None:
+        super().__init__()
+        self.time = time
+        self.unit = unit
+        self._name: TableOptionKey = TableOptionKey.TARGET_LAG
+
+    @staticmethod
+    def create(
+        value: Union["TargetLagOption", Tuple[int, TimeUnit], KeywordOptionType]
+    ) -> Optional[TableOption]:
+        if isinstance(value, NoneType):
+            return value
+
+        if isinstance(value, Tuple):
+            time, unit = value
+            value = TargetLagOption(time, unit)
+
+        if isinstance(value, TargetLagOption):
+            return value
+
+        if isinstance(value, (KeywordOption, SnowflakeKeyword)):
+            return KeywordOption.create(TableOptionKey.TARGET_LAG, value)
+
+        return TableOption._get_invalid_table_option(
+            TableOptionKey.TARGET_LAG,
+            str(type(value).__name__),
+            [
+                TargetLagOption.__name__,
+                f"Tuple[int, {TimeUnit.__name__}])",
+                SnowflakeKeyword.__name__,
+            ],
+        )
+
+    def __get_expression(self):
+        return f"'{str(self.time)} {str(self.unit.value)}'"
+
+    @property
+    def priority(self) -> Priority:
+        return Priority.HIGH
+
+    def _render(self, compiler) -> str:
+        return self.template() % (self.__get_expression())
+
+    def __repr__(self) -> str:
+        return "TargetLagOption(%s)" % self.__get_expression()
+
+
+TargetLagOptionType = Union[TargetLagOption, Tuple[int, TimeUnit]]

snowflake/sqlalchemy/sql/custom_schema/snowflake_table.py

@@ -0,0 +1,70 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Any
+
+from sqlalchemy.sql.schema import MetaData, SchemaItem
+
+from .table_from_query import TableFromQueryBase
+
+
+class SnowflakeTable(TableFromQueryBase):
+    """
+    A class representing a table in Snowflake with configurable options and settings.
+
+    While it does not support reflection at this time, it provides a flexible
+    interface for creating tables and management.
+
+    For further information on this clause, please refer to: https://docs.snowflake.com/en/sql-reference/sql/create-table
+    Example usage:
+
+    SnowflakeTable(
+        table_name,
+        metadata,
+        Column("id", Integer, primary_key=True),
+        Column("name", String),
+        cluster_by = ["id", text("name > 5")]
+    )
+
+    Example using explict options:
+
+        SnowflakeTable(
+        table_name,
+        metadata,
+        Column("id", Integer, primary_key=True),
+        Column("name", String),
+        cluster_by = ClusterByOption("id", text("name > 5"))
+    )
+
+    """
+
+    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 "SnowflakeTable(%s)" % ", ".join(
+            [repr(self.name)]
+            + [repr(self.metadata)]
+            + [repr(x) for x in self.columns]
+            + [repr(self.cluster_by)]
+            + [repr(self.as_query)]
+            + [f"{k}={repr(getattr(self, k))}" for k in ["schema"]]
+        )

snowflake/sqlalchemy/sql/custom_schema/table_from_query.py

@@ -0,0 +1,54 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+import typing
+from typing import Any, Optional
+
+from sqlalchemy.sql import Selectable
+from sqlalchemy.sql.schema import Column, MetaData, SchemaItem
+
+from .clustered_table import ClusteredTableBase
+from .options.as_query_option import AsQueryOption, AsQueryOptionType
+from .options.table_option import TableOptionKey
+
+
+class TableFromQueryBase(ClusteredTableBase):
+
+    @property
+    def as_query(self) -> Optional[AsQueryOption]:
+        return self._get_dialect_option(TableOptionKey.AS_QUERY)
+
+    def __init__(
+        self,
+        name: str,
+        metadata: MetaData,
+        *args: SchemaItem,
+        as_query: AsQueryOptionType = None,
+        **kw: Any,
+    ) -> None:
+        items = [item for item in args]
+        as_query = AsQueryOption.create(as_query)  # noqa
+        kw.update(self._as_dialect_options([as_query]))
+        if (
+            isinstance(as_query, AsQueryOption)
+            and isinstance(as_query.query, Selectable)
+            and not self.__has_defined_columns(items)
+        ):
+            columns = self.__create_columns_from_selectable(as_query.query)
+            args = items + columns
+        super().__init__(name, metadata, *args, **kw)
+
+    def __has_defined_columns(self, items: typing.List[SchemaItem]) -> bool:
+        for item in items:
+            if isinstance(item, Column):
+                return True
+
+    def __create_columns_from_selectable(
+        self, selectable: Selectable
+    ) -> Optional[typing.List[Column]]:
+        if not isinstance(selectable, Selectable):
+            return
+        columns: typing.List[Column] = []
+        for _, c in selectable.exported_columns.items():
+            columns += [Column(c.name, c.type)]
+        return columns