dagster-snowflake 0.21.5__py3-none-any.whl → 0.25.8__py3-none-any.whl

dagster_snowflake/__init__.py

@@ -1,15 +1,16 @@
 from dagster._core.libraries import DagsterLibraryRegistry
 
-from .ops import snowflake_op_for_query as snowflake_op_for_query
-from .resources import (
+from dagster_snowflake.ops import snowflake_op_for_query as snowflake_op_for_query
+from dagster_snowflake.resources import (
     SnowflakeConnection as SnowflakeConnection,
     SnowflakeResource as SnowflakeResource,
+    fetch_last_updated_timestamps as fetch_last_updated_timestamps,
     snowflake_resource as snowflake_resource,
 )
-from .snowflake_io_manager import (
+from dagster_snowflake.snowflake_io_manager import (
     SnowflakeIOManager as SnowflakeIOManager,
     build_snowflake_io_manager as build_snowflake_io_manager,
 )
-from .version import __version__
+from dagster_snowflake.version import __version__
 
 DagsterLibraryRegistry.register("dagster-snowflake", __version__)

dagster_snowflake/constants.py

@@ -0,0 +1,7 @@
+# Description: This file contains the Snowflake connection identifiers for the Snowflake partner account.
+# The connection identifiers are used to identify the partner account when connecting to Snowflake.
+
+# We use different connection identifiers for different connection code paths to ensure that each is
+# working as expected.
+SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER = "DagsterLabs_Dagster"
+SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER_SQLALCHEMY = "DagsterLabs_Dagster_SqlAlchemy"

dagster_snowflake/ops.py

@@ -4,6 +4,7 @@ from dagster import (
     op,
 )
 from dagster._core.definitions.input import In
+from dagster._core.storage.tags import COMPUTE_KIND_TAG
 
 
 def _core_create_snowflake_command(dagster_decorator, decorator_name, sql, parameters=None):
@@ -14,7 +15,7 @@ def _core_create_snowflake_command(dagster_decorator, decorator_name, sql, param
         name=f"snowflake_{decorator_name}",
         ins={"start": In(Nothing)},
         required_resource_keys={"snowflake"},
-        tags={"kind": "sql", "sql": sql},
+        tags={COMPUTE_KIND_TAG: "sql", "sql": sql},
     )
     def snowflake_fn(context):
         context.resources.snowflake.execute_query(sql=sql, parameters=parameters)

dagster_snowflake/py.typed

@@ -0,0 +1 @@
+partial

dagster_snowflake/resources.py

@@ -2,6 +2,7 @@ import base64
 import sys
 import warnings
 from contextlib import closing, contextmanager
+from datetime import datetime
 from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Union
 
 import dagster._check as check
@@ -14,11 +15,15 @@ from dagster import (
     resource,
 )
 from dagster._annotations import public
-from dagster._config.pythonic_config.pydantic_compat_layer import compat_model_validator
 from dagster._core.definitions.resource_definition import dagster_maintained_resource
 from dagster._core.storage.event_log.sql_event_log import SqlDbConnection
 from dagster._utils.cached_method import cached_method
-from pydantic import Field, validator
+from pydantic import Field, model_validator, validator
+
+from dagster_snowflake.constants import (
+    SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER,
+    SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER_SQLALCHEMY,
+)
 
 try:
     import snowflake.connector
@@ -229,7 +234,7 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
             "Indicate alternative database connection engine. Permissible option is "
             "'sqlalchemy' otherwise defaults to use the Snowflake Connector for Python."
         ),
-        is_required=False,
+        is_required=False,  # type: ignore
     )
 
     cache_column_metadata: Optional[str] = Field(
@@ -253,6 +258,15 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
         default=None,
         description="Optional parameter to specify the authentication mechanism to use.",
     )
+    additional_snowflake_connection_args: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description=(
+            "Additional keyword arguments to pass to the snowflake.connector.connect function. For a full list of"
+            " available arguments, see the `Snowflake documentation"
+            " <https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect>`__."
+            " This config will be ignored if using the sqlalchemy connector."
+        ),
+    )
 
     @validator("paramstyle")
     def validate_paramstyle(cls, v: Optional[str]) -> Optional[str]:
@@ -272,7 +286,7 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
             )
         return v
 
-    @compat_model_validator(mode="before")
+    @model_validator(mode="before")
     def validate_authentication(cls, values):
         auths_set = 0
         auths_set += 1 if values.get("password") is not None else 0
@@ -332,6 +346,10 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
         ):
             conn_args["private_key"] = self._snowflake_private_key(self._resolved_config_dict)
 
+        conn_args["application"] = SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER
+
+        if self._resolved_config_dict.get("additional_snowflake_connection_args") is not None:
+            conn_args.update(self._resolved_config_dict["additional_snowflake_connection_args"])
         return conn_args
 
     @property
@@ -352,6 +370,7 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
             )
             if self._resolved_config_dict.get(k) is not None
         }
+        conn_args["application"] = SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER_SQLALCHEMY
 
         return conn_args
 
@@ -713,3 +732,63 @@ def snowflake_resource(context) -> SnowflakeConnection:
     return SnowflakeConnection(
         config=context, log=context.log, snowflake_connection_resource=snowflake_resource
     )
+
+
+def fetch_last_updated_timestamps(
+    *,
+    snowflake_connection: Union[SqlDbConnection, snowflake.connector.SnowflakeConnection],
+    schema: str,
+    tables: Sequence[str],
+    database: Optional[str] = None,
+    ignore_missing_tables: Optional[bool] = False,
+) -> Mapping[str, datetime]:
+    """Fetch the last updated times of a list of tables in Snowflake.
+
+    If the underlying query to fetch the last updated time returns no results, a ValueError will be raised.
+
+    Args:
+        snowflake_connection (Union[SqlDbConnection, SnowflakeConnection]): A connection to Snowflake.
+            Accepts either a SnowflakeConnection or a sqlalchemy connection object,
+            which are the two types of connections emittable from the snowflake resource.
+        schema (str): The schema of the tables to fetch the last updated time for.
+        tables (Sequence[str]): A list of table names to fetch the last updated time for.
+        database (Optional[str]): The database of the table. Only required if the connection
+            has not been set with a database.
+        ignore_missing_tables (Optional[bool]): If True, tables not found in Snowflake
+            will be excluded from the result.
+
+    Returns:
+        Mapping[str, datetime]: A dictionary of table names to their last updated time in UTC.
+    """
+    check.invariant(len(tables) > 0, "Must provide at least one table name to query upon.")
+    # Table names in snowflake's information schema are stored in uppercase
+    uppercase_tables = [table.upper() for table in tables]
+    tables_str = ", ".join([f"'{table_name}'" for table_name in uppercase_tables])
+    fully_qualified_table_name = (
+        f"{database}.information_schema.tables" if database else "information_schema.tables"
+    )
+
+    query = f"""
+    SELECT table_name, CONVERT_TIMEZONE('UTC', last_altered) AS last_altered
+    FROM {fully_qualified_table_name}
+    WHERE table_schema = '{schema}' AND table_name IN ({tables_str});
+    """
+    result = snowflake_connection.cursor().execute(query)
+    if not result:
+        raise ValueError("No results returned from Snowflake update time query.")
+
+    result_mapping = {table_name: last_altered for table_name, last_altered in result}
+    result_correct_case = {}
+    for table_name in tables:
+        if table_name.upper() not in result_mapping:
+            if ignore_missing_tables:
+                continue
+            raise ValueError(f"Table {table_name} could not be found.")
+        last_altered = result_mapping[table_name.upper()]
+        check.invariant(
+            isinstance(last_altered, datetime),
+            "Expected last_altered to be a datetime, but it was not.",
+        )
+        result_correct_case[table_name] = last_altered
+
+    return result_correct_case

dagster_snowflake/snowflake_io_manager.py

@@ -1,11 +1,9 @@
 from abc import abstractmethod
 from contextlib import contextmanager
-from typing import Optional, Sequence, Type, cast
+from typing import Any, Dict, Optional, Sequence, Type, cast
 
 from dagster import IOManagerDefinition, OutputContext, io_manager
-from dagster._config.pythonic_config import (
-    ConfigurableIOManagerFactory,
-)
+from dagster._config.pythonic_config import ConfigurableIOManagerFactory
 from dagster._core.definitions.time_window_partitions import TimeWindow
 from dagster._core.storage.db_io_manager import (
     DbClient,
@@ -16,9 +14,9 @@ from dagster._core.storage.db_io_manager import (
 )
 from dagster._core.storage.io_manager import dagster_maintained_io_manager
 from pydantic import Field
-from sqlalchemy.exc import ProgrammingError
+from snowflake.connector.errors import ProgrammingError
 
-from .resources import SnowflakeResource
+from dagster_snowflake.resources import SnowflakeResource
 
 SNOWFLAKE_DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
 
@@ -46,15 +44,22 @@ def build_snowflake_io_manager(
             from dagster import Definitions
 
             @asset(
-                key_prefix=["my_schema"]  # will be used as the schema in snowflake
+                key_prefix=["my_prefix"]
+                metadata={"schema": "my_schema"} # will be used as the schema in snowflake
             )
             def my_table() -> pd.DataFrame:  # the name of the asset will be the table name
                 ...
 
+            @asset(
+                key_prefix=["my_schema"]  # will be used as the schema in snowflake
+            )
+            def my_second_table() -> pd.DataFrame:  # the name of the asset will be the table name
+                ...
+
             snowflake_io_manager = build_snowflake_io_manager([SnowflakePandasTypeHandler(), SnowflakePySparkTypeHandler()])
 
             defs = Definitions(
-                assets=[my_table],
+                assets=[my_table, my_second_table],
                 resources={
                     "io_manager": snowflake_io_manager.configured({
                         "database": "my_database",
@@ -64,12 +69,38 @@ def build_snowflake_io_manager(
                 }
             )
 
-        If you do not provide a schema, Dagster will determine a schema based on the assets and ops using
-        the IO Manager. For assets, the schema will be determined from the asset key,
-        as shown in the above example. The final prefix before the asset name will be used as the schema. For example,
-        if the asset ``my_table`` had the key prefix ``["snowflake", "my_schema"]``, the schema ``my_schema`` will be
-        used. For ops, the schema can be specified by including a ``schema`` entry in output metadata. If ``schema`` is not provided
-        via config or on the asset/op, ``public`` will be used for the schema.
+        You can set a default schema to store the assets using the ``schema`` configuration value of the Snowflake I/O
+        Manager. This schema will be used if no other schema is specified directly on an asset or op.
+
+        .. code-block:: python
+
+            defs = Definitions(
+                assets=[my_table]
+                resources={"io_manager" snowflake_io_manager.configured(
+                    {"database": "my_database", "schema": "my_schema", ...} # will be used as the schema
+                )}
+            )
+
+
+        On individual assets, you an also specify the schema where they should be stored using metadata or
+        by adding a ``key_prefix`` to the asset key. If both ``key_prefix`` and metadata are defined, the metadata will
+        take precedence.
+
+        .. code-block:: python
+
+            @asset(
+                key_prefix=["my_schema"]  # will be used as the schema in snowflake
+            )
+            def my_table() -> pd.DataFrame:
+                ...
+
+            @asset(
+                metadata={"schema": "my_schema"}  # will be used as the schema in snowflake
+            )
+            def my_other_table() -> pd.DataFrame:
+                ...
+
+        For ops, the schema can be specified by including a "schema" entry in output metadata.
 
         .. code-block:: python
 
@@ -77,9 +108,10 @@ def build_snowflake_io_manager(
                 out={"my_table": Out(metadata={"schema": "my_schema"})}
             )
             def make_my_table() -> pd.DataFrame:
-                # the returned value will be stored at my_schema.my_table
                 ...
 
+        If none of these is provided, the schema will default to "public".
+
         To only use specific columns of a table as input to a downstream op or asset, add the metadata ``columns`` to the
         In or AssetIn.
 
@@ -134,16 +166,42 @@ class SnowflakeIOManager(ConfigurableIOManagerFactory):
             defs = Definitions(
                 assets=[my_table],
                 resources={
-                    "io_manager": MySnowflakeIOManager(database="MY_DATABASE", account=EnvVar("SNOWFLAKE_ACCOUNT"), ...)
+                    "io_manager": MySnowflakeIOManager(database="my_database", account=EnvVar("SNOWFLAKE_ACCOUNT"), ...)
                 }
             )
 
-        If you do not provide a schema, Dagster will determine a schema based on the assets and ops using
-        the IO Manager. For assets, the schema will be determined from the asset key,
-        as shown in the above example. The final prefix before the asset name will be used as the schema. For example,
-        if the asset ``my_table`` had the key prefix ``["snowflake", "my_schema"]``, the schema ``my_schema`` will be
-        used. For ops, the schema can be specified by including a ``schema`` entry in output metadata. If ``schema`` is not provided
-        via config or on the asset/op, ``public`` will be used for the schema.
+        You can set a default schema to store the assets using the ``schema`` configuration value of the Snowflake I/O
+        Manager. This schema will be used if no other schema is specified directly on an asset or op.
+
+        .. code-block:: python
+
+            defs = Definitions(
+                assets=[my_table]
+                resources={
+                    "io_manager" MySnowflakeIOManager(database="my_database", schema="my_schema", ...)
+                }
+            )
+
+
+        On individual assets, you an also specify the schema where they should be stored using metadata or
+        by adding a ``key_prefix`` to the asset key. If both ``key_prefix`` and metadata are defined, the metadata will
+        take precedence.
+
+        .. code-block:: python
+
+            @asset(
+                key_prefix=["my_schema"]  # will be used as the schema in snowflake
+            )
+            def my_table() -> pd.DataFrame:
+                ...
+
+            @asset(
+                metadata={"schema": "my_schema"}  # will be used as the schema in snowflake
+            )
+            def my_other_table() -> pd.DataFrame:
+                ...
+
+        For ops, the schema can be specified by including a "schema" entry in output metadata.
 
         .. code-block:: python
 
@@ -151,9 +209,10 @@ class SnowflakeIOManager(ConfigurableIOManagerFactory):
                 out={"my_table": Out(metadata={"schema": "my_schema"})}
             )
             def make_my_table() -> pd.DataFrame:
-                # the returned value will be stored at my_schema.my_table
                 ...
 
+        If none of these is provided, the schema will default to "public".
+
         To only use specific columns of a table as input to a downstream op or asset, add the metadata ``columns`` to the
         In or AssetIn.
 
@@ -220,6 +279,15 @@ class SnowflakeIOManager(ConfigurableIOManagerFactory):
         default=None,
         description="Optional parameter to specify the authentication mechanism to use.",
     )
+    additional_snowflake_connection_args: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description=(
+            "Additional keyword arguments to pass to the snowflake.connector.connect function. For a full list of"
+            " available arguments, see the `Snowflake documentation"
+            " <https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect>`__."
+            " This config will be ignored if using the sqlalchemy connector."
+        ),
+    )
 
     @staticmethod
     @abstractmethod
@@ -288,26 +356,31 @@ class SnowflakeDbClient(DbClient):
             if context.resource_config
             else {}
         )
-        with SnowflakeResource(
-            schema=table_slice.schema, connector="sqlalchemy", **no_schema_config
-        ).get_connection(raw_conn=False) as conn:
+        with SnowflakeResource(schema=table_slice.schema, **no_schema_config).get_connection(  # pyright: ignore[reportArgumentType]
+            raw_conn=False
+        ) as conn:
             yield conn
 
     @staticmethod
     def ensure_schema_exists(context: OutputContext, table_slice: TableSlice, connection) -> None:
-        schemas = connection.execute(
-            f"show schemas like '{table_slice.schema}' in database {table_slice.database}"
-        ).fetchall()
+        schemas = (
+            connection.cursor()
+            .execute(f"show schemas like '{table_slice.schema}' in database {table_slice.database}")
+            .fetchall()
+        )
         if len(schemas) == 0:
-            connection.execute(f"create schema {table_slice.schema};")
+            connection.cursor().execute(f"create schema {table_slice.schema};")
 
     @staticmethod
     def delete_table_slice(context: OutputContext, table_slice: TableSlice, connection) -> None:
         try:
-            connection.execute(_get_cleanup_statement(table_slice))
-        except ProgrammingError:
-            # table doesn't exist yet, so ignore the error
-            pass
+            connection.cursor().execute(_get_cleanup_statement(table_slice))
+        except ProgrammingError as e:
+            if "does not exist" in e.msg:  # type: ignore
+                # table doesn't exist yet, so ignore the error
+                return
+            else:
+                raise
 
     @staticmethod
     def get_select_statement(table_slice: TableSlice) -> str:

dagster_snowflake/version.py

@@ -1 +1 @@
-__version__ = "0.21.5"
+__version__ = "0.25.8"

{dagster_snowflake-0.21.5.dist-info → dagster_snowflake-0.25.8.dist-info}/METADATA

@@ -1,23 +1,25 @@
 Metadata-Version: 2.1
 Name: dagster-snowflake
-Version: 0.21.5
+Version: 0.25.8
 Summary: Package for Snowflake Dagster framework components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-snowflake
 Author: Dagster Labs
 Author-email: hello@dagsterlabs.com
 License: Apache-2.0
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9,<3.13
 License-File: LICENSE
-Requires-Dist: dagster ==1.5.5
-Requires-Dist: snowflake-connector-python >=2.1.0
+Requires-Dist: dagster ==1.9.8
+Requires-Dist: snowflake-connector-python >=3.4.0
+Requires-Dist: pyOpenSSL >=22.1.0
 Provides-Extra: pandas
 Requires-Dist: pandas ; extra == 'pandas'
-Requires-Dist: snowflake-connector-python[pandas] >=2.1.0 ; extra == 'pandas'
+Requires-Dist: snowflake-connector-python[pandas] >=3.4.0 ; extra == 'pandas'
 Provides-Extra: snowflake.sqlalchemy
 Requires-Dist: sqlalchemy !=1.4.42 ; extra == 'snowflake.sqlalchemy'
 Requires-Dist: snowflake-sqlalchemy ; extra == 'snowflake.sqlalchemy'

dagster_snowflake-0.25.8.dist-info/RECORD

@@ -0,0 +1,12 @@
+dagster_snowflake/__init__.py,sha256=MeLU-ghhmyMk82I2b3qCz9h7mxgb2eMxXfkqGBjI29k,673
+dagster_snowflake/constants.py,sha256=26PGL1eFncm0WfbgZjA7jaznfaRx75vVJnqbieVYeco,481
+dagster_snowflake/ops.py,sha256=L_MP28fLm7_hrJmzMoDocLwyvVnkpy1LVwUSLIArKWc,2225
+dagster_snowflake/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
+dagster_snowflake/resources.py,sha256=_ERXAGne5q_Z0JG92fBvPOnAySUtWiqJmjU1qYddZwI,32084
+dagster_snowflake/snowflake_io_manager.py,sha256=d_fFr7Fk95uNQ6PNvbxCYC6dgH2yyzJBAx_KVC8HtTQ,17763
+dagster_snowflake/version.py,sha256=bkf--LdvzFbtaf_d7GSQn5A2t-yDZ-SX33wKaaeqHsY,23
+dagster_snowflake-0.25.8.dist-info/LICENSE,sha256=TMatHW4_G9ldRdodEAp-l2Xa2WvsdeOh60E3v1R2jis,11349
+dagster_snowflake-0.25.8.dist-info/METADATA,sha256=pUY1pTPFxlJFwf1fYWvBztWMjEUJeJmHJ6UM7qEH8gY,1103
+dagster_snowflake-0.25.8.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+dagster_snowflake-0.25.8.dist-info/top_level.txt,sha256=uECYCiluOxLQ996SCUPBBwdK0CTyz45FjWqf7WDqMMc,18
+dagster_snowflake-0.25.8.dist-info/RECORD,,

dagster_snowflake-0.21.5.dist-info/RECORD

@@ -1,10 +0,0 @@
-dagster_snowflake/__init__.py,sha256=svCkKOLfCh3_1DBoTtrEa8Hu7puNB4c7JX4FvIGbUf0,537
-dagster_snowflake/ops.py,sha256=cxCMvbzsELa62ax3AWodKSpmad3Kb6NLcZlj6SW2KVg,2159
-dagster_snowflake/resources.py,sha256=ATW__nAcep_t7muqylM6Ux5jxVCgSM9uz7tpjpLx-jU,28337
-dagster_snowflake/snowflake_io_manager.py,sha256=C5sHxtsHpuKsRWdbsCv9hVmZ7mvProSCgwoQ1UeCy6Y,15478
-dagster_snowflake/version.py,sha256=yq-fGfdU069XbMVu8cXj7x9Ga-YNsSY_1SmDswa-WEI,23
-dagster_snowflake-0.21.5.dist-info/LICENSE,sha256=TMatHW4_G9ldRdodEAp-l2Xa2WvsdeOh60E3v1R2jis,11349
-dagster_snowflake-0.21.5.dist-info/METADATA,sha256=Fw_rmHZNz0yVkFZUIChtL7IBSybN6rdtHJdyCE9V-Gg,1039
-dagster_snowflake-0.21.5.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-dagster_snowflake-0.21.5.dist-info/top_level.txt,sha256=uECYCiluOxLQ996SCUPBBwdK0CTyz45FjWqf7WDqMMc,18
-dagster_snowflake-0.21.5.dist-info/RECORD,,