UV package management and Snowflake logging extras #35

dominictarro · 2024-09-20T20:16:59Z

Two-part PR

Transferred package management to uv.
Added a Snowflake logging method to an extras module.

Snowflake Logging

Gives prefect_snowflake.SnowflakeConnector users the ability to log commands executed by the cursor if it fails or during debug. This is particularly useful for pipelines generating SQL on the fly.

from prefecto.blocks import lazy_load
from prefecto.ext.snowflake import PrefectoSnowflakeConnector


class Blocks:
    SNOWFLAKE_BLOCK_NAME: str = "snowflake-block-name"

    @property
    @lazy_load("SNOWFLAKE_BLOCK_NAME")
    def snowflake(self) -> PrefectoSnowflakeConnector:
        """Snowflake connector block."""

blocks = Blocks()

with blocks.snowflake.get_connection().cursor() as cursor:
    cursor.execute(
        "SELECT * FROM table WHERE id = %(id)s AND secret = %(secret)s",
        params={"id": 123, "secret": "shhh"},
        obfuscate_params=["secret"],
        command_id="secret-selection",
    )

INFO - [secret-selection] Beginning command.
DEBUG - [secret-selection] Executing command:
SELECT * FROM table WHERE id = 123 AND secret = ****
INFO - [secret-selection] Command executed successfully.

- update project def - in workflows, install uv and install packages through it (cache slower than reinstall) - init uv.lock with existing deps

- logs start/end of snowflake command - optionally display or mask parameters in command - display command on failure - extended connector block for wrapping the execute method

ybryzTE

Looks good. Got couple questions and propostions before approving.

ybryzTE · 2024-09-24T16:24:07Z

src/prefecto/ext/snowflake.py

+    except Exception as e:
+        if __python_version__.major > 3 or (
+            __python_version__.major == 3 and __python_version__.minor >= 11
+        ):
+            e.add_note(f"Command failed\n{formatted_command_string}")
+        logger.error("Command failed.", exc_info=e, stack_info=True)
+        raise e


I think we can change it to this to avoid checking the versions:

if hasattr(e, "add_note") and callable(e.add_note): e.add_note("Command failed."...)

ybryzTE · 2024-09-24T16:45:59Z

src/prefecto/ext/snowflake.py

+        [Snowflake Python connector documentation](https://docs.snowflake.com/en/user-guide/python-connector-api.html#connect).
+
+        Args:
+            **connect_kwargs: Additional arguments to pass to
+                `snowflake.connector.connect`.
+
+        Returns:
+            An authenticated Snowflake connection.
+
+        Example:
+            Get Snowflake connection with only block configuration:
+            ```python
+            from prefecto.ext.snowflake import PrefectoSnowflakeCredentials as SnowflakeCredentials
+
+            snowflake_credentials_block = SnowflakeCredentials.load("BLOCK_NAME")
+
+            connection = snowflake_credentials_block.get_client()
+            ```
+
+            Get Snowflake connector scoped to a specified database:
+            ```python
+            from prefecto.ext.snowflake import PrefectoSnowflakeCredentials as SnowflakeCredentials
+
+            snowflake_credentials_block = SnowflakeCredentials.load("BLOCK_NAME")
+
+            connection = snowflake_credentials_block.get_client(database="my_database")
+            ```
+        """  # noqa
+        connect_params = {
+            # required to track task's usage in the Snowflake Partner Network Portal
+            "application": "Prefect_Snowflake_Collection",
+            **self.credentials.model_dump(
+                exclude_unset=True, exclude={"block_type_slug"}
+            ),
+            **connect_kwargs,
+        }
+
+        for key, value in connect_params.items():
+            if isinstance(value, (SecretStr, SecretBytes)):
+                connect_params[key] = connect_params[key].get_secret_value()
+
+        # set authenticator to the actual okta_endpoint
+        if connect_params.get("authenticator") == "okta_endpoint":
+            endpoint = connect_params.pop("endpoint", None) or connect_params.pop(
+                "okta_endpoint", None
+            )  # okta_endpoint is deprecated
+            connect_params["authenticator"] = endpoint
+
+        private_der_key = self.credentials.resolve_private_key()
+        if private_der_key is not None:
+            connect_params["private_key"] = private_der_key
+            connect_params.pop("password", None)
+            connect_params.pop("private_key_passphrase", None)
+
+        return PrefectoSnowflakeConnection(**connect_params)
+
+    def get_connection(self, **connect_kwargs: Any) -> PrefectoSnowflakeConnection:
+        """
+        Returns an authenticated connection that can be
+        used to command from Snowflake databases.
+
+        Args:
+            **connect_kwargs: Additional arguments to pass to
+                `snowflake.connector.connect`.
+
+        Returns:
+            The authenticated SnowflakeConnection.
+
+        Examples:
+            ```python
+            from prefecto.ext.snowflake import PrefectoSnowflakeCredentials as SnowflakeCredentials
+            from prefecto.ext.snowflake import PrefectoSnowflakeConnector as SnowflakeConnector
+
+            snowflake_credentials = SnowflakeCredentials(
+                account="account",
+                user="user",
+                password="password",
+            )
+            snowflake_connector = SnowflakeConnector(
+                database="database",
+                warehouse="warehouse",
+                schema="schema",
+                credentials=snowflake_credentials
+            )
+            with snowflake_connector.get_connection() as connection:
+                ...
+            ```
+        """
+        if self._connection is not None:
+            return self._connection
+
+        connect_params = {
+            "database": self.database,
+            "warehouse": self.warehouse,
+            "schema": self.schema_,
+        }
+        connection = self._credentials_get_client(**connect_kwargs, **connect_params)
+        self._connection = connection
+        self.logger.info("Started a new connection to Snowflake.")
+        return connection


Correct me if I'm wrong, but this one follows the same logic found in prefect_snowflake, just injecting the custom PrefectoSnowflakeConnection at the end of _credentials_get_client, whereas in prefect_snowflake it calls snowflake.connector.connect(**connect_params) instead. Are there other mechanisms that prefect_snowflake provides to accomplish this without rewriting the same logic? I'd rather use PrefectoSnowflakeCursor directly with the call to cursor(PrefectoSnowflakeCursor) on the original connection object returned from SnowflakeConnector. Otherwise, one possible approach that I can think of to remove the duplicate logic is to use composition over inheritance.

I like that idea more. I removed the relationship to prefect-snowflake and focused on keeping this a cursor class like DictCursor. It raises the question of "why is this in prefecto?", but it will be used alongside Prefect enough to be worth it as an extension.

from snowflake.connector import connect from prefecto.ext.snowflake import LogCursor with connect(...).cursor(LogCursor) as cursor: cursor.execute( "SELECT * FROM table WHERE id = %(id)s AND secret = %(secret)s", params={"id": 123, "secret": "shhh"}, obfuscate_params=["secret"], command_id="secret-selection", )

INFO - [secret-selection] Beginning command. DEBUG - [secret-selection] Executing command: SELECT * FROM table WHERE id = 123 AND secret = **** INFO - [secret-selection] Command executed successfully.

src/prefecto/ext/snowflake.py

This approach emphasizes just using the Snowflake connection object's `.cursor(cursor_cls)` functionality.

ybryzTE

Looks good!

dominictarro added 5 commits September 19, 2024 15:34

refactor: pivot to uv package manager

a614b80

- update project def - in workflows, install uv and install packages through it (cache slower than reinstall) - init uv.lock with existing deps

refactor: remove docstring requirement

598dc87

feat: command execution logging for snowflake connector block

4bd74f8

- logs start/end of snowflake command - optionally display or mask parameters in command - display command on failure - extended connector block for wrapping the execute method

refactor: migrate extension modules to a higher extensions module

821d1e7

docs: Touch up dosctrings and add the mkdocs page

7b81c26

dominictarro added the enhancement New feature or request label Sep 20, 2024

dominictarro self-assigned this Sep 20, 2024

fix: formatted all files with pre-commit

5d78028

dominictarro requested a review from ybryzTE September 20, 2024 20:21

ybryzTE reviewed Sep 24, 2024

View reviewed changes

dominictarro added 6 commits September 24, 2024 15:29

refactor: simplify error handling logic in _execute

2a34dad

refactor: Remove the block implementations and update docstrings

97ed027

This approach emphasizes just using the Snowflake connection object's `.cursor(cursor_cls)` functionality.

conf: update dependency requirements

8cf52e7

refactor: Move snowflake custom cursor unit test to extensions test dir

a4707a8

docs: Add LogCursor class docstring

2c674fb

docs: Fix snowflake ext module docstring

6bd5c3a

ybryzTE approved these changes Sep 25, 2024

View reviewed changes

dominictarro merged commit 0607705 into main Sep 25, 2024
4 checks passed

dominictarro deleted the feat-uv-and-snowflake-logging-extra branch September 25, 2024 14:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

UV package management and Snowflake logging extras #35

UV package management and Snowflake logging extras #35

dominictarro commented Sep 20, 2024

ybryzTE left a comment

ybryzTE Sep 24, 2024

ybryzTE Sep 24, 2024

dominictarro Sep 25, 2024

ybryzTE left a comment

UV package management and Snowflake logging extras #35

UV package management and Snowflake logging extras #35

Conversation

dominictarro commented Sep 20, 2024

Snowflake Logging

ybryzTE left a comment

Choose a reason for hiding this comment

ybryzTE Sep 24, 2024

Choose a reason for hiding this comment

ybryzTE Sep 24, 2024

Choose a reason for hiding this comment

dominictarro Sep 25, 2024

Choose a reason for hiding this comment

ybryzTE left a comment

Choose a reason for hiding this comment