snowflake-sqlalchemy 1.5.2__py3-none-any.whl

snowflake/sqlalchemy/custom_commands.py

@@ -0,0 +1,621 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from collections.abc import Sequence
+from typing import List
+
+from sqlalchemy import false, true
+from sqlalchemy.sql.ddl import DDLElement
+from sqlalchemy.sql.dml import UpdateBase
+from sqlalchemy.sql.elements import ClauseElement
+from sqlalchemy.sql.roles import FromClauseRole
+from sqlalchemy.util.compat import string_types
+
+NoneType = type(None)
+
+
+def translate_bool(bln):
+    if bln:
+        return true()
+    return false()
+
+
+class MergeInto(UpdateBase):
+    __visit_name__ = "merge_into"
+    _bind = None
+
+    def __init__(self, target, source, on):
+        self.target = target
+        self.source = source
+        self.on = on
+        self.clauses = []
+
+    class clause(ClauseElement):
+        __visit_name__ = "merge_into_clause"
+
+        def __init__(self, command):
+            self.set = {}
+            self.predicate = None
+            self.command = command
+
+        def __repr__(self):
+            case_predicate = (
+                f" AND {str(self.predicate)}" if self.predicate is not None else ""
+            )
+            if self.command == "INSERT":
+                sets, sets_tos = zip(*self.set.items())
+                return "WHEN NOT MATCHED{} THEN {} ({}) VALUES ({})".format(
+                    case_predicate,
+                    self.command,
+                    ", ".join(sets),
+                    ", ".join(map(str, sets_tos)),
+                )
+            else:
+                # WHEN MATCHED clause
+                sets = (
+                    ", ".join([f"{set[0]} = {set[1]}" for set in self.set.items()])
+                    if self.set
+                    else ""
+                )
+                return "WHEN MATCHED{} THEN {}{}".format(
+                    case_predicate,
+                    self.command,
+                    f" SET {str(sets)}" if self.set else "",
+                )
+
+        def values(self, **kwargs):
+            self.set = kwargs
+            return self
+
+        def where(self, expr):
+            self.predicate = expr
+            return self
+
+    def __repr__(self):
+        clauses = " ".join([repr(clause) for clause in self.clauses])
+        return f"MERGE INTO {self.target} USING {self.source} ON {self.on}" + (
+            f" {clauses}" if clauses else ""
+        )
+
+    def when_matched_then_update(self):
+        clause = self.clause("UPDATE")
+        self.clauses.append(clause)
+        return clause
+
+    def when_matched_then_delete(self):
+        clause = self.clause("DELETE")
+        self.clauses.append(clause)
+        return clause
+
+    def when_not_matched_then_insert(self):
+        clause = self.clause("INSERT")
+        self.clauses.append(clause)
+        return clause
+
+
+class FilesOption:
+    """
+    Class to represent FILES option for the snowflake COPY INTO statement
+    """
+
+    def __init__(self, file_names: List[str]):
+        self.file_names = file_names
+
+    def __str__(self):
+        the_files = ["'" + f.replace("'", "\\'") + "'" for f in self.file_names]
+        return f"({','.join(the_files)})"
+
+
+class CopyInto(UpdateBase):
+    """Copy Into Command base class, for documentation see:
+    https://docs.snowflake.net/manuals/sql-reference/sql/copy-into-location.html"""
+
+    __visit_name__ = "copy_into"
+    _bind = None
+
+    def __init__(self, from_, into, formatter=None):
+        self.from_ = from_
+        self.into = into
+        self.formatter = formatter
+        self.copy_options = {}
+
+    def __repr__(self):
+        """
+        repr for debugging / logging purposes only. For compilation logic, see
+        the corresponding visitor in base.py
+        """
+        return f"COPY INTO {self.into} FROM {repr(self.from_)} {repr(self.formatter)} ({self.copy_options})"
+
+    def bind(self):
+        return None
+
+    def force(self, force):
+        if not isinstance(force, bool):
+            raise TypeError("Parameter force should be a boolean value")
+        self.copy_options.update({"FORCE": translate_bool(force)})
+        return self
+
+    def single(self, single_file):
+        if not isinstance(single_file, bool):
+            raise TypeError("Parameter single_file should  be a boolean value")
+        self.copy_options.update({"SINGLE": translate_bool(single_file)})
+        return self
+
+    def maxfilesize(self, max_size):
+        if not isinstance(max_size, int):
+            raise TypeError("Parameter max_size should be an integer value")
+        self.copy_options.update({"MAX_FILE_SIZE": max_size})
+        return self
+
+    def files(self, file_names):
+        self.copy_options.update({"FILES": FilesOption(file_names)})
+        return self
+
+    def pattern(self, pattern):
+        self.copy_options.update({"PATTERN": pattern})
+        return self
+
+
+class CopyFormatter(ClauseElement):
+    """
+    Base class for Formatter specifications inside a COPY INTO statement. May also
+    be used to create a named format.
+    """
+
+    __visit_name__ = "copy_formatter"
+
+    def __init__(self, format_name=None):
+        self.options = dict()
+        if format_name:
+            self.options["format_name"] = format_name
+
+    def __repr__(self):
+        """
+        repr for debugging / logging purposes only. For compilation logic, see
+        the corresponding visitor in base.py
+        """
+        return f"FILE_FORMAT=({self.options})"
+
+    @staticmethod
+    def value_repr(name, value):
+        """
+        Make a SQL-suitable representation of "value". This is called from
+        the corresponding visitor function (base.py/visit_copy_formatter())
+        - in case of a format name: return it without quotes
+        - in case of a string: enclose in quotes: "value"
+        - in case of a tuple of length 1: enclose the only element in brackets: (value)
+            Standard stringification of Python would append a trailing comma: (value,)
+            which is not correct in SQL
+        - otherwise: just convert to str as is: value
+        """
+        if name == "format_name":
+            return value
+        elif isinstance(value, str):
+            return f"'{value}'"
+        elif isinstance(value, tuple) and len(value) == 1:
+            return f"('{value[0]}')"
+        else:
+            return str(value)
+
+
+class CSVFormatter(CopyFormatter):
+    file_format = "csv"
+
+    def compression(self, comp_type):
+        """String (constant) that specifies to compresses the unloaded data files using the specified compression algorithm."""
+        if isinstance(comp_type, string_types):
+            comp_type = comp_type.lower()
+        _available_options = [
+            "auto",
+            "gzip",
+            "bz2",
+            "brotli",
+            "zstd",
+            "deflate",
+            "raw_deflate",
+            None,
+        ]
+        if comp_type not in _available_options:
+            raise TypeError(f"Compression type should be one of : {_available_options}")
+        self.options["COMPRESSION"] = comp_type
+        return self
+
+    def _check_delimiter(self, delimiter, delimiter_txt):
+        """
+        Check if a delimiter is either a string of length 1 or an integer. In case of
+        a string delimiter, take into account that the actual string may be longer,
+        but still evaluate to a single character (like "\\n" or r"\n"
+        """
+        if isinstance(delimiter, NoneType):
+            return
+        if isinstance(delimiter, string_types):
+            delimiter_processed = delimiter.encode().decode("unicode_escape")
+            if len(delimiter_processed) == 1:
+                return
+        if isinstance(delimiter, int):
+            return
+        raise TypeError(
+            f"{delimiter_txt} should be a single character, that is either a string, or a number"
+        )
+
+    def record_delimiter(self, deli_type):
+        """Character that separates records in an unloaded file."""
+        self._check_delimiter(deli_type, "Record delimiter")
+        if isinstance(deli_type, int):
+            self.options["RECORD_DELIMITER"] = hex(deli_type)
+        else:
+            self.options["RECORD_DELIMITER"] = deli_type
+        return self
+
+    def field_delimiter(self, deli_type):
+        """Character that separates fields in an unloaded file."""
+        self._check_delimiter(deli_type, "Field delimiter")
+        if isinstance(deli_type, int):
+            self.options["FIELD_DELIMITER"] = hex(deli_type)
+        else:
+            self.options["FIELD_DELIMITER"] = deli_type
+        return self
+
+    def file_extension(self, ext):
+        """String that specifies the extension for files unloaded to a stage. Accepts any extension. The user is
+        responsible for specifying a valid file extension that can be read by the desired software or service.
+        """
+        if not isinstance(ext, (NoneType, string_types)):
+            raise TypeError("File extension should be a string")
+        self.options["FILE_EXTENSION"] = ext
+        return self
+
+    def date_format(self, dt_frmt):
+        """String that defines the format of date values in the unloaded data files."""
+        if not isinstance(dt_frmt, string_types):
+            raise TypeError("Date format should be a string")
+        self.options["DATE_FORMAT"] = dt_frmt
+        return self
+
+    def time_format(self, tm_frmt):
+        """String that defines the format of time values in the unloaded data files."""
+        if not isinstance(tm_frmt, string_types):
+            raise TypeError("Time format should be a string")
+        self.options["TIME_FORMAT"] = tm_frmt
+        return self
+
+    def timestamp_format(self, tmstmp_frmt):
+        """String that defines the format of timestamp values in the unloaded data files."""
+        if not isinstance(tmstmp_frmt, string_types):
+            raise TypeError("Timestamp format should be a string")
+        self.options["TIMESTAMP_FORMAT"] = tmstmp_frmt
+        return self
+
+    def binary_format(self, bin_fmt):
+        """Character used as the escape character for any field values. The option can be used when unloading data
+        from binary columns in a table."""
+        if isinstance(bin_fmt, string_types):
+            bin_fmt = bin_fmt.lower()
+        _available_options = ["hex", "base64", "utf8"]
+        if bin_fmt not in _available_options:
+            raise TypeError(f"Binary format should be one of : {_available_options}")
+        self.options["BINARY_FORMAT"] = bin_fmt
+        return self
+
+    def escape(self, esc):
+        """Character used as the escape character for any field values."""
+        self._check_delimiter(esc, "Escape")
+        if isinstance(esc, int):
+            self.options["ESCAPE"] = hex(esc)
+        else:
+            self.options["ESCAPE"] = esc
+        return self
+
+    def escape_unenclosed_field(self, esc):
+        """Single character string used as the escape character for unenclosed field values only."""
+        self._check_delimiter(esc, "Escape unenclosed field")
+        if isinstance(esc, int):
+            self.options["ESCAPE_UNENCLOSED_FIELD"] = hex(esc)
+        else:
+            self.options["ESCAPE_UNENCLOSED_FIELD"] = esc
+        return self
+
+    def field_optionally_enclosed_by(self, enc):
+        """Character used to enclose strings. Either None, ', or \"."""
+        _available_options = [None, "'", '"']
+        if enc not in _available_options:
+            raise TypeError(f"Enclosing string should be one of : {_available_options}")
+        self.options["FIELD_OPTIONALLY_ENCLOSED_BY"] = enc
+        return self
+
+    def null_if(self, null_value):
+        """Copying into a table these strings will be replaced by a NULL, while copying out of Snowflake will replace
+        NULL values with the first string"""
+        if not isinstance(null_value, Sequence):
+            raise TypeError("Parameter null_value should be an iterable")
+        self.options["NULL_IF"] = tuple(null_value)
+        return self
+
+    def skip_header(self, skip_header):
+        """
+        Number of header rows to be skipped at the beginning of the file
+        """
+        if not isinstance(skip_header, int):
+            raise TypeError("skip_header  should be an int")
+        self.options["SKIP_HEADER"] = skip_header
+        return self
+
+    def trim_space(self, trim_space):
+        """
+        Remove leading or trailing white spaces
+        """
+        if not isinstance(trim_space, bool):
+            raise TypeError("trim_space should be a bool")
+        self.options["TRIM_SPACE"] = trim_space
+        return self
+
+    def error_on_column_count_mismatch(self, error_on_col_count_mismatch):
+        """
+        Generate a parsing error if the number of delimited columns (i.e. fields) in
+        an input data file does not match the number of columns in the corresponding table.
+        """
+        if not isinstance(error_on_col_count_mismatch, bool):
+            raise TypeError("skip_header  should be a bool")
+        self.options["ERROR_ON_COLUMN_COUNT_MISMATCH"] = error_on_col_count_mismatch
+        return self
+
+
+class JSONFormatter(CopyFormatter):
+    """Format specific functions"""
+
+    file_format = "json"
+
+    def compression(self, comp_type):
+        """String (constant) that specifies to compresses the unloaded data files using the specified compression algorithm."""
+        if isinstance(comp_type, string_types):
+            comp_type = comp_type.lower()
+        _available_options = [
+            "auto",
+            "gzip",
+            "bz2",
+            "brotli",
+            "zstd",
+            "deflate",
+            "raw_deflate",
+            None,
+        ]
+        if comp_type not in _available_options:
+            raise TypeError(f"Compression type should be one of : {_available_options}")
+        self.options["COMPRESSION"] = comp_type
+        return self
+
+    def file_extension(self, ext):
+        """String that specifies the extension for files unloaded to a stage. Accepts any extension. The user is
+        responsible for specifying a valid file extension that can be read by the desired software or service.
+        """
+        if not isinstance(ext, (NoneType, string_types)):
+            raise TypeError("File extension should be a string")
+        self.options["FILE_EXTENSION"] = ext
+        return self
+
+
+class PARQUETFormatter(CopyFormatter):
+    """Format specific functions"""
+
+    file_format = "parquet"
+
+    def snappy_compression(self, comp):
+        """Enable, or disable snappy compression"""
+        if not isinstance(comp, bool):
+            raise TypeError("Comp should be a Boolean value")
+        self.options["SNAPPY_COMPRESSION"] = translate_bool(comp)
+        return self
+
+    def compression(self, comp):
+        """
+        Set compression type
+        """
+        if not isinstance(comp, str):
+            raise TypeError("Comp should be a str value")
+        self.options["COMPRESSION"] = comp
+        return self
+
+    def binary_as_text(self, value):
+        """Enable, or disable binary as text"""
+        if not isinstance(value, bool):
+            raise TypeError("binary_as_text should be a Boolean value")
+        self.options["BINARY_AS_TEXT"] = translate_bool(value)
+        return self
+
+
+class ExternalStage(ClauseElement, FromClauseRole):
+    """External Stage descriptor"""
+
+    __visit_name__ = "external_stage"
+    _hide_froms = ()
+
+    @staticmethod
+    def prepare_namespace(namespace):
+        return f"{namespace}." if not namespace.endswith(".") else namespace
+
+    @staticmethod
+    def prepare_path(path):
+        return f"/{path}" if not path.startswith("/") else path
+
+    def __init__(self, name, path=None, namespace=None, file_format=None):
+        self.name = name
+        self.path = self.prepare_path(path) if path else ""
+        self.namespace = self.prepare_namespace(namespace) if namespace else ""
+        self.file_format = file_format
+
+    def __repr__(self):
+        return f"@{self.namespace}{self.name}{self.path} ({self.file_format})"
+
+    @classmethod
+    def from_parent_stage(cls, parent_stage, path, file_format=None):
+        """
+        Extend an existing parent stage (with or without path) with an
+        additional sub-path
+        """
+        return cls(
+            parent_stage.name,
+            f"{parent_stage.path}/{path}",
+            parent_stage.namespace,
+            file_format,
+        )
+
+
+class CreateFileFormat(DDLElement):
+    """
+    Encapsulates a CREATE FILE FORMAT statement; using a format description (as in
+    a COPY INTO statement) and a format name.
+    """
+
+    __visit_name__ = "create_file_format"
+
+    def __init__(self, format_name, formatter, replace_if_exists=False):
+        super().__init__()
+        self.format_name = format_name
+        self.formatter = formatter
+        self.replace_if_exists = replace_if_exists
+
+
+class CreateStage(DDLElement):
+    """
+    Encapsulates a CREATE STAGE statement, using a container (physical base for the
+    stage) and the actual ExternalStage object.
+    """
+
+    __visit_name__ = "create_stage"
+
+    def __init__(self, container, stage, replace_if_exists=False, *, temporary=False):
+        super().__init__()
+        self.container = container
+        self.temporary = temporary
+        self.stage = stage
+        self.replace_if_exists = replace_if_exists
+
+
+class AWSBucket(ClauseElement):
+    """AWS S3 bucket descriptor"""
+
+    __visit_name__ = "aws_bucket"
+
+    def __init__(self, bucket, path=None):
+        self.bucket = bucket
+        self.path = path
+        self.encryption_used = {}
+        self.credentials_used = {}
+
+    @classmethod
+    def from_uri(cls, uri):
+        if uri[0:5] != "s3://":
+            raise ValueError(f"Invalid AWS bucket URI: {uri}")
+        b = uri[5:].split("/", 1)
+        if len(b) == 1:
+            bucket, path = b[0], None
+        else:
+            bucket, path = b
+        return cls(bucket, path)
+
+    def __repr__(self):
+        credentials = "CREDENTIALS=({})".format(
+            " ".join(f"{n}='{v}'" for n, v in self.credentials_used.items())
+        )
+        encryption = "ENCRYPTION=({})".format(
+            " ".join(
+                f"{n}='{v}'" if isinstance(v, string_types) else f"{n}={v}"
+                for n, v in self.encryption_used.items()
+            )
+        )
+        uri = "'s3://{}{}'".format(self.bucket, f"/{self.path}" if self.path else "")
+        return "{}{}{}".format(
+            uri,
+            f" {credentials}" if self.credentials_used else "",
+            f" {encryption}" if self.encryption_used else "",
+        )
+
+    def credentials(
+        self, aws_role=None, aws_key_id=None, aws_secret_key=None, aws_token=None
+    ):
+        if aws_role is None and (aws_key_id is None and aws_secret_key is None):
+            raise ValueError(
+                "Either 'aws_role', or aws_key_id and aws_secret_key has to be supplied"
+            )
+        if aws_role:
+            self.credentials_used = {"AWS_ROLE": aws_role}
+        else:
+            self.credentials_used = {
+                "AWS_SECRET_KEY": aws_secret_key,
+                "AWS_KEY_ID": aws_key_id,
+            }
+            if aws_token:
+                self.credentials_used["AWS_TOKEN"] = aws_token
+        return self
+
+    def encryption_aws_cse(self, master_key):
+        self.encryption_used = {"TYPE": "AWS_CSE", "MASTER_KEY": master_key}
+        return self
+
+    def encryption_aws_sse_s3(self):
+        self.encryption_used = {"TYPE": "AWS_SSE_S3"}
+        return self
+
+    def encryption_aws_sse_kms(self, kms_key_id=None):
+        self.encryption_used = {"TYPE": "AWS_SSE_KMS"}
+        if kms_key_id:
+            self.encryption_used["KMS_KEY_ID"] = kms_key_id
+        return self
+
+
+class AzureContainer(ClauseElement):
+    """Microsoft Azure Container descriptor"""
+
+    __visit_name__ = "azure_container"
+
+    def __init__(self, account, container, path=None):
+        self.account = account
+        self.container = container
+        self.path = path
+        self.encryption_used = {}
+        self.credentials_used = {}
+
+    @classmethod
+    def from_uri(cls, uri):
+        if uri[0:8] != "azure://":
+            raise ValueError(f"Invalid Azure Container URI: {uri}")
+        account, uri = uri[8:].split(".", 1)
+        if uri[0:22] != "blob.core.windows.net/":
+            raise ValueError(f"Invalid Azure Container URI: {uri}")
+        b = uri[22:].split("/", 1)
+        if len(b) == 1:
+            container, path = b[0], None
+        else:
+            container, path = b
+        return cls(account, container, path)
+
+    def __repr__(self):
+        credentials = "CREDENTIALS=({})".format(
+            " ".join(f"{n}='{v}'" for n, v in self.credentials_used.items())
+        )
+        encryption = "ENCRYPTION=({})".format(
+            " ".join(
+                f"{n}='{v}'" if isinstance(v, string_types) else f"{n}={v}"
+                for n, v in self.encryption_used.items()
+            )
+        )
+        uri = "'azure://{}.blob.core.windows.net/{}{}'".format(
+            self.account, self.container, f"/{self.path}" if self.path else ""
+        )
+        return "{}{}{}".format(
+            uri,
+            f" {credentials}" if self.credentials_used else "",
+            f" {encryption}" if self.encryption_used else "",
+        )
+
+    def credentials(self, azure_sas_token):
+        self.credentials_used = {"AZURE_SAS_TOKEN": azure_sas_token}
+        return self
+
+    def encryption_azure_cse(self, master_key):
+        self.encryption_used = {"TYPE": "AZURE_CSE", "MASTER_KEY": master_key}
+        return self
+
+
+CopyIntoStorage = CopyInto