soda-snowflake 4.0.5__tar.gz

soda_snowflake-4.0.5/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: soda-snowflake
+Version: 4.0.5
+Requires-Dist: soda-core==4.0.5
+Requires-Dist: snowflake-connector-python>=3.0
+Dynamic: requires-dist

soda_snowflake-4.0.5/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

soda_snowflake-4.0.5/setup.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+
+from setuptools import setup
+
+package_name = "soda-snowflake"
+package_version = "4.0.5"
+description = "Soda Snowflake V4"
+
+requires = [
+    f"soda-core=={package_version}",
+    "snowflake-connector-python>=3.0",
+]  # 4.0 is also fine, but for backwards compatibility we include 3.0
+
+setup(
+    name=package_name,
+    version=package_version,
+    install_requires=requires,
+    package_dir={"": "src"},
+    entry_points={
+        "soda.plugins.data_source.snowflake": [
+            "SnowflakeDataSourceImpl = soda_snowflake.common.data_sources.snowflake_data_source:SnowflakeDataSourceImpl",
+        ],
+    },
+)

soda_snowflake-4.0.5/src/soda_snowflake/common/data_sources/snowflake_data_source.py

@@ -0,0 +1,200 @@
+from logging import Logger
+from numbers import Number
+from typing import TYPE_CHECKING, Optional
+
+from soda_core.common.data_source_connection import DataSourceConnection
+from soda_core.common.data_source_impl import DataSourceImpl
+from soda_core.common.logging_constants import soda_logger
+from soda_core.common.metadata_types import SodaDataTypeName
+from soda_core.common.soda_cloud_dto import SamplerType
+from soda_core.common.sql_ast import COLUMN, COUNT, DISTINCT, TUPLE, VALUES
+from soda_core.common.sql_dialect import SqlDialect
+from soda_core.contracts.impl.contract_verification_impl import ContractImpl
+from soda_snowflake.common.data_sources.snowflake_data_source_connection import (
+    SnowflakeDataSource as SnowflakeDataSourceModel,
+)
+from soda_snowflake.common.data_sources.snowflake_data_source_connection import (
+    SnowflakeDataSourceConnection,
+)
+
+logger: Logger = soda_logger
+
+if TYPE_CHECKING:
+    from soda_core.contracts.impl.contract_verification_impl import ContractImpl
+
+TIMESTAMP_WITHOUT_TIME_ZONE = "timestamp without time zone"
+TIMESTAMP_WITH_TIME_ZONE = "timestamp with time zone"
+TIMESTAMP_WITH_LOCAL_TIME_ZONE = "timestamp with local time zone"
+
+
+class SnowflakeDataSourceImpl(DataSourceImpl, model_class=SnowflakeDataSourceModel):
+    def __init__(self, data_source_model: SnowflakeDataSourceModel, connection: Optional[DataSourceConnection] = None):
+        super().__init__(data_source_model=data_source_model, connection=connection)
+
+    def _create_sql_dialect(self) -> SqlDialect:
+        return SnowflakeSqlDialect(data_source_impl=self)
+
+    def _create_data_source_connection(self) -> DataSourceConnection:
+        return SnowflakeDataSourceConnection(
+            name=self.data_source_model.name, connection_properties=self.data_source_model.connection_properties
+        )
+
+    def switch_warehouse(self, warehouse: str, contract_impl: ContractImpl) -> None:
+        if warehouse and contract_impl.datasource_warehouse != warehouse:
+            if contract_impl.datasource_warehouse is None:
+                logger.info(
+                    f"Setting warehouse to '{warehouse}' for Contract verification of dataset '{contract_impl.dataset_identifier.to_string()}'"
+                )
+            else:
+                logger.info(
+                    f"Switching warehouse from '{contract_impl.datasource_warehouse}' to '{warehouse}' for Contract verification of dataset '{contract_impl.dataset_identifier.to_string()}'"
+                )
+            self._execute_switch_warehouse(warehouse)
+        else:
+            logger.info(
+                f"Using warehouse '{contract_impl.datasource_warehouse}' for Contract verification of dataset '{contract_impl.dataset_identifier.to_string()}'"
+            )
+
+    def _execute_switch_warehouse(self, warehouse: str) -> None:
+        switch_warehouse_sql = f"USE WAREHOUSE {warehouse}"
+        self.execute_query(switch_warehouse_sql)
+
+    def get_current_warehouse(self) -> Optional[str]:
+        sql = "SELECT CURRENT_WAREHOUSE()"
+        result = self.execute_query(sql)
+        result_rows = result.rows
+        row = result_rows[0] if result_rows else None
+
+        return row[0] if row and row[0] else None
+
+
+class SnowflakeSqlDialect(SqlDialect):
+    SODA_DATA_TYPE_SYNONYMS = (
+        (SodaDataTypeName.TEXT, SodaDataTypeName.VARCHAR, SodaDataTypeName.CHAR),
+        (
+            SodaDataTypeName.NUMERIC,
+            SodaDataTypeName.DECIMAL,
+            SodaDataTypeName.INTEGER,
+            SodaDataTypeName.BIGINT,
+            SodaDataTypeName.SMALLINT,
+        ),
+        (SodaDataTypeName.FLOAT, SodaDataTypeName.DOUBLE),
+    )
+
+    def default_casify(self, identifier: str) -> str:
+        return identifier.upper()
+
+    def metadata_casify(self, identifier: str) -> str:
+        """Metadata identifiers are not uppercased for Snowflake."""
+        return identifier
+
+    def create_table_casify_qualified_name(self, qualified_name: str) -> str:
+        # Parse the last element from the qualified name and casify it
+        if "." in qualified_name:  # It is an actual qualified name, not just a table name
+            name_parts = qualified_name.split(".")
+            # Caseify the last element
+            name_parts[-1] = name_parts[-1].upper()
+            return ".".join(name_parts)
+        # It is not a fully qualified name, just a table name
+        return qualified_name.upper()
+
+    def build_cte_values_sql(self, values: VALUES, alias_columns: list[COLUMN] | None) -> str:
+        return " SELECT * FROM VALUES\n" + ",\n".join([self.build_expression_sql(value) for value in values.values])
+
+    def _build_tuple_sql(self, tuple: TUPLE) -> str:
+        if tuple.check_context(COUNT) and tuple.check_context(DISTINCT):
+            return self._build_tuple_sql_in_distinct(tuple)
+        return f"{super()._build_tuple_sql(tuple)}"
+
+    def _build_tuple_sql_in_distinct(self, tuple: TUPLE) -> str:
+        return f"ARRAY_CONSTRUCT{super()._build_tuple_sql(tuple)}"
+
+    def _get_data_type_name_synonyms(self) -> list[list[str]]:
+        # Implements data type synonyms
+        # Each list should represent a list of synonyms
+        return [
+            ["varchar", "text", "string"],
+            ["number", "decimal", "numeric", "int", "integer", "bigint", "smallint", "tinyint", "byteint"],
+            ["float", "float4", "float8", "double", "double precision", "real"],
+            ["timestamp", "datetime", "timestamp_ntz", TIMESTAMP_WITHOUT_TIME_ZONE],
+            ["timestamp_ltz", TIMESTAMP_WITH_LOCAL_TIME_ZONE],
+            ["timestamp_tz", TIMESTAMP_WITH_TIME_ZONE],
+        ]
+
+    def get_data_source_data_type_name_by_soda_data_type_names(self) -> dict:
+        """
+        Maps DBDataType names to data source type names.
+        """
+        return {
+            SodaDataTypeName.CHAR: "char",  # Note: by default a char is 1 byte in Snowflake!
+            SodaDataTypeName.VARCHAR: "varchar",
+            SodaDataTypeName.TEXT: "text",  # alias for varchar
+            SodaDataTypeName.SMALLINT: "smallint",
+            SodaDataTypeName.INTEGER: "integer",
+            SodaDataTypeName.BIGINT: "bigint",
+            SodaDataTypeName.DECIMAL: "number",  # decimal & numeric → number
+            SodaDataTypeName.NUMERIC: "number",
+            SodaDataTypeName.FLOAT: "float",  # float / double → float
+            SodaDataTypeName.DOUBLE: "float",
+            SodaDataTypeName.TIMESTAMP: "timestamp_ntz",  # default timestamp in snowflake
+            SodaDataTypeName.TIMESTAMP_TZ: "timestamp_tz",
+            SodaDataTypeName.DATE: "date",
+            SodaDataTypeName.TIME: "time",
+            SodaDataTypeName.BOOLEAN: "boolean",
+        }
+
+    def data_type_has_parameter_character_maximum_length(self, data_type_name) -> bool:
+        return data_type_name.lower() in ["varchar", "char", "character", "text"]
+
+    # TODO: test this thorough. The code here is generated using AI just to be able to test the E2E.
+    def get_soda_data_type_name_by_data_source_data_type_names(self) -> dict[str, SodaDataTypeName]:
+        return {
+            "varchar": SodaDataTypeName.VARCHAR,
+            "char": SodaDataTypeName.CHAR,
+            "character": SodaDataTypeName.CHAR,
+            "text": SodaDataTypeName.TEXT,
+            "string": SodaDataTypeName.VARCHAR,
+            "smallint": SodaDataTypeName.SMALLINT,
+            "integer": SodaDataTypeName.INTEGER,
+            "int": SodaDataTypeName.INTEGER,
+            "bigint": SodaDataTypeName.BIGINT,
+            "tinyint": SodaDataTypeName.SMALLINT,
+            "byteint": SodaDataTypeName.SMALLINT,
+            "number": SodaDataTypeName.NUMERIC,
+            "decimal": SodaDataTypeName.DECIMAL,
+            "numeric": SodaDataTypeName.NUMERIC,
+            "float": SodaDataTypeName.FLOAT,
+            "float4": SodaDataTypeName.FLOAT,
+            "float8": SodaDataTypeName.DOUBLE,
+            "double": SodaDataTypeName.DOUBLE,
+            "double precision": SodaDataTypeName.DOUBLE,
+            "real": SodaDataTypeName.FLOAT,
+            "timestamp": SodaDataTypeName.TIMESTAMP,
+            "timestamp_ntz": SodaDataTypeName.TIMESTAMP,
+            TIMESTAMP_WITHOUT_TIME_ZONE: SodaDataTypeName.TIMESTAMP,
+            "datetime": SodaDataTypeName.TIMESTAMP,
+            "timestamp_tz": SodaDataTypeName.TIMESTAMP_TZ,
+            TIMESTAMP_WITH_TIME_ZONE: SodaDataTypeName.TIMESTAMP_TZ,
+            "timestamp_ltz": SodaDataTypeName.TIMESTAMP_TZ,
+            TIMESTAMP_WITH_LOCAL_TIME_ZONE: SodaDataTypeName.TIMESTAMP_TZ,
+            "date": SodaDataTypeName.DATE,
+            "time": SodaDataTypeName.TIME,
+            "boolean": SodaDataTypeName.BOOLEAN,
+        }
+
+    def data_type_has_parameter_datetime_precision(self, data_type_name) -> bool:
+        return data_type_name.lower() in [
+            "timestamp",
+            "timestamp_ntz",
+            TIMESTAMP_WITHOUT_TIME_ZONE,
+            "timestamp_tz",
+            TIMESTAMP_WITH_TIME_ZONE,
+            "timestamp_ltz",
+            TIMESTAMP_WITH_LOCAL_TIME_ZONE,
+        ]
+
+    def _build_sample_sql(self, sampler_type: str, sample_size: Number) -> str:
+        if sampler_type == SamplerType.ABSOLUTE_LIMIT:
+            return f"TABLESAMPLE ({int(sample_size)} ROWS)"
+        else:
+            raise ValueError(f"Unsupported sample type: {sampler_type}")

soda_snowflake-4.0.5/src/soda_snowflake/common/data_sources/snowflake_data_source_connection.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC
+from pathlib import Path
+from typing import Dict, Literal, Optional
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from pydantic import Field, SecretStr, field_validator
+from snowflake import connector
+from soda_core.common.data_source_connection import DataSourceConnection
+from soda_core.common.logging_constants import soda_logger
+from soda_core.model.data_source.data_source import DataSourceBase
+from soda_core.model.data_source.data_source_connection_properties import (
+    DataSourceConnectionProperties,
+)
+
+logger: logging.Logger = soda_logger
+
+USER_DESCRIPTION = "Username for authentication"
+
+
+class SnowflakeConnectionProperties(DataSourceConnectionProperties, ABC):
+    ...
+
+
+class SnowflakeSharedConnectionProperties(SnowflakeConnectionProperties, ABC):
+    account: str = Field(..., description="Snowflake account identifier")
+    warehouse: Optional[str] = Field(None, description="Name of the warehouse to use")
+    database: Optional[str] = Field(None, description="Name of the database to use")
+    role: Optional[str] = Field(None, description="Role to assume after connecting")
+    session_parameters: Optional[Dict[str, str]] = Field(None, description="Session-level parameters")
+    host: Optional[str] = Field(None, description="Host name of the Snowflake account")
+
+
+class SnowflakePasswordAuth(SnowflakeSharedConnectionProperties):
+    user: str = Field(..., description=USER_DESCRIPTION)
+    password: SecretStr = Field(..., description="User password")
+
+
+class SnowflakeKeyPairAuth(SnowflakeSharedConnectionProperties):
+    user: str = Field(..., description=USER_DESCRIPTION)
+    private_key: SecretStr = Field(..., description="Private key for authentication")
+    private_key_passphrase: Optional[SecretStr] = Field(None, description="Passphrase if private key is encrypted")
+
+    def to_connection_kwargs(self) -> dict:
+        connection_kwargs = super().to_connection_kwargs()
+        connection_kwargs["private_key"] = self._decrypt(self.private_key, self.private_key_passphrase)
+        return connection_kwargs
+
+    def _decrypt(self, private_key: str, private_key_passphrase: Optional[SecretStr]) -> bytes:
+        private_key_bytes = private_key.get_secret_value().encode()
+        private_key_passphrase_bytes = (
+            private_key_passphrase.get_secret_value().encode() if private_key_passphrase else None
+        )
+
+        p_key = serialization.load_pem_private_key(
+            private_key_bytes, password=private_key_passphrase_bytes, backend=default_backend()
+        )
+
+        return p_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+
+class SnowflakeKeyPairFileAuth(SnowflakeSharedConnectionProperties):
+    user: str = Field(..., description=USER_DESCRIPTION)
+    private_key_path: Path = Field(..., description="Path to private key file")
+    private_key_passphrase: Optional[SecretStr] = Field(None, description="Passphrase if private key is encrypted")
+
+    def to_connection_kwargs(self) -> dict:
+        connection_kwargs = super().to_connection_kwargs()
+        connection_kwargs["private_key_file"] = self.private_key_path
+        if self.private_key_passphrase is not None:
+            pwd = self.private_key_passphrase.get_secret_value()
+            if pwd:
+                connection_kwargs["private_key_file_pwd"] = pwd
+        return connection_kwargs
+
+
+class SnowflakeOAuthAuth(SnowflakeSharedConnectionProperties):
+    authenticator: Literal["oauth"] = Field(..., description="Use OAuth access token")
+    token: SecretStr = Field(..., description="OAuth access token")
+
+
+class SnowflakeClientCredentialsOAuthAuth(SnowflakeSharedConnectionProperties):
+    # User is not required for OAuth Client Credentials Flow.
+    authenticator: Literal["OAUTH_CLIENT_CREDENTIALS"] = Field(
+        ..., description="Authenticator to use for OAuth Client Credentials Flow."
+    )
+    oauth_client_id: SecretStr = Field(..., description="Client ID for OAuth Client Credentials Flow.")
+    oauth_client_secret: SecretStr = Field(..., description="Client secret for OAuth Client Credentials Flow.")
+    oauth_token_request_url: SecretStr = Field(..., description="Token request URL for OAuth Client Credentials Flow.")
+    oauth_scope: Optional[SecretStr] = Field(None, description="Scope for OAuth Client Credentials Flow if required.")
+
+
+class SnowflakeSSOAuth(SnowflakeSharedConnectionProperties):
+    # User is not required for SSO Flow.
+    authenticator: Literal["externalbrowser"] = Field("externalbrowser", description="Use external browser SSO login")
+
+
+class SnowflakeDataSource(DataSourceBase, ABC):
+    type: Literal["snowflake"] = Field("snowflake")
+    connection_properties: SnowflakeConnectionProperties = Field(
+        ..., alias="connection", description="Snowflake connection configuration"
+    )
+
+    @field_validator("connection_properties", mode="before")
+    @classmethod
+    def infer_connection_type(cls, value):
+        if "password" in value:
+            return SnowflakePasswordAuth(**value)
+        elif "private_key" in value:
+            return SnowflakeKeyPairAuth(**value)
+        elif "private_key_path" in value:
+            return SnowflakeKeyPairFileAuth(**value)
+        elif "token" in value:
+            return SnowflakeOAuthAuth(**value)
+        elif value.get("authenticator") == "externalbrowser":
+            return SnowflakeSSOAuth(**value)
+        elif value.get("authenticator") == "OAUTH_CLIENT_CREDENTIALS":
+            return SnowflakeClientCredentialsOAuthAuth(**value)
+        raise ValueError("Could not infer Snowflake connection type from input")
+
+
+class SnowflakeDataSourceConnection(DataSourceConnection):
+    def __init__(self, name: str, connection_properties: DataSourceConnectionProperties):
+        super().__init__(name, connection_properties)
+
+    def _create_connection(
+        self,
+        config: SnowflakeConnectionProperties,
+    ):
+        return connector.connect(
+            application="Soda",
+            **config.to_connection_kwargs(),
+        )

soda_snowflake-4.0.5/src/soda_snowflake/test_helpers/snowflake_data_source_test_helper.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import os
+
+from helpers.data_source_test_helper import DataSourceTestHelper
+
+
+class SnowflakeDataSourceTestHelper(DataSourceTestHelper):
+    def _create_database_name(self) -> str | None:
+        return os.getenv("SNOWFLAKE_DATABASE", "soda_test")
+
+    def _create_data_source_yaml_str(self) -> str:
+        """
+        Called in _create_data_source_impl to initialized self.data_source_impl
+        self.database_name and self.schema_name are available if appropriate for the data source type
+        """
+        return f"""
+            type: snowflake
+            name: {self.name}
+            connection:
+                account: {os.getenv("SNOWFLAKE_ACCOUNT")}
+                user: {os.getenv("SNOWFLAKE_USER")}
+                password: {os.getenv("SNOWFLAKE_PASSWORD")}
+                database: {self.dataset_prefix[0]}
+        """
+
+    def _adjust_schema_name(self, schema_name: str) -> str:
+        return schema_name.upper()

soda_snowflake-4.0.5/src/soda_snowflake.egg-info/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: soda-snowflake
+Version: 4.0.5
+Requires-Dist: soda-core==4.0.5
+Requires-Dist: snowflake-connector-python>=3.0
+Dynamic: requires-dist

soda_snowflake-4.0.5/src/soda_snowflake.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+setup.py
+src/soda_snowflake.egg-info/PKG-INFO
+src/soda_snowflake.egg-info/SOURCES.txt
+src/soda_snowflake.egg-info/dependency_links.txt
+src/soda_snowflake.egg-info/entry_points.txt
+src/soda_snowflake.egg-info/requires.txt
+src/soda_snowflake.egg-info/top_level.txt
+src/soda_snowflake/common/data_sources/snowflake_data_source.py
+src/soda_snowflake/common/data_sources/snowflake_data_source_connection.py
+src/soda_snowflake/test_helpers/snowflake_data_source_test_helper.py

soda_snowflake-4.0.5/src/soda_snowflake.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

soda_snowflake-4.0.5/src/soda_snowflake.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[soda.plugins.data_source.snowflake]
+SnowflakeDataSourceImpl = soda_snowflake.common.data_sources.snowflake_data_source:SnowflakeDataSourceImpl

soda_snowflake-4.0.5/src/soda_snowflake.egg-info/requires.txt

@@ -0,0 +1,2 @@
+soda-core==4.0.5
+snowflake-connector-python>=3.0

soda_snowflake-4.0.5/src/soda_snowflake.egg-info/top_level.txt

@@ -0,0 +1 @@
+soda_snowflake