detquantlib 5.4.0__tar.gz → 5.5.0__tar.gz

{detquantlib-5.4.0 → detquantlib-5.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: detquantlib
-Version: 5.4.0
+Version: 5.5.0
 Summary: An internal library containing functions and classes that can be used across Quant models.
 License-File: LICENSE.txt
 Author: Dynamic Energy Trading

{detquantlib-5.4.0 → detquantlib-5.5.0}/detquantlib/data/databases/snowflake_client.py

@@ -4,24 +4,22 @@ Client for working with a Snowflake database from Python.
 This module provides a configuration dataclass and a client class that wrap
 ``snowflake-connector-python`` with a small, safe API for interacting with Snowflake. Data is
 exchanged as pandas DataFrames.
-
-Example:
-    >>> with SnowflakeClient(SnowflakeConfig.from_env()) as client:
-    ...     df = client.query("SELECT CURRENT_VERSION() AS version")
 """
 
-from __future__ import annotations
-
+# Python built-in packages
 import logging
 import os
 import re
 from dataclasses import dataclass, field
+from types import TracebackType
 from typing import Any, Iterator, Optional, Sequence, Union
 
+# Third-party packages
 import pandas as pd
 import snowflake.connector
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.errors import Error as SnowflakeError
+from snowflake.connector.pandas_tools import write_pandas
 
 logger = logging.getLogger(__name__)
 
@@ -108,10 +106,27 @@ class SnowflakeClient:
     Attributes:
         config: The connection settings used by this client.
 
-    Example:
-        >>> with SnowflakeClient(SnowflakeConfig.from_env()) as client:
-        ...     df = client.query("SELECT * FROM orders WHERE status = %s", ("SHIPPED",))
+    Examples:
+        Example 1: Loading a table
+        >>> with SnowflakeClient() as client:
         ...     customers = client.load_table("CUSTOMERS", columns=["ID", "NAME"], limit=100)
+
+        Example 2: Standard query method
+        >>> with SnowflakeClient() as client:
+        ...     df = client.query("SELECT * FROM orders WHERE status = %s", ("SHIPPED",))
+
+        Example 3: Batch query method to stream a large table in chunks
+        >>> with SnowflakeClient(SnowflakeConfig.from_env()) as client:
+        ...     for chunk in client.query_batches("SELECT * FROM BIG_TABLE"):
+        ...         process(chunk)
+
+        Example 4: Writing a DataFrame to a table, creating the table if needed
+        >>> with SnowflakeClient() as client:
+        ...     client.write_table(df, "DB.SCHEMA.ORDERS", auto_create_table=True)
+
+        Example 5: Running a statement that does not return data
+        >>> with SnowflakeClient() as client:
+        ...     client.execute("DELETE FROM ORDERS WHERE STATUS = %s", ("CANCELLED",))
     """
 
     def __init__(self, config: Optional[SnowflakeConfig] = None) -> None:
@@ -228,7 +243,12 @@ class SnowflakeClient:
         self.connect()
         return self
 
-    def __exit__(self, exc_type, exc, tb) -> None:
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
         """
         Closes the connection when leaving a ``with`` block.
 
@@ -359,6 +379,120 @@ class SnowflakeClient:
             cols = [d[0] for d in cur.description]
             return pd.DataFrame(cur.fetchall(), columns=cols)
 
+    # ---------------------------------------------------------------------------------------------
+    # Writing data
+    # ---------------------------------------------------------------------------------------------
+    def write_table(
+        self,
+        df: pd.DataFrame,
+        table: str,
+        overwrite: bool = False,
+        auto_create_table: bool = False,
+        chunk_size: Optional[int] = None,
+    ) -> int:
+        """
+        Writes a DataFrame to a table.
+
+        Data is uploaded in bulk with the connector's ``write_pandas`` (staged Parquet files loaded
+        with ``COPY INTO``), which is much faster than row-by-row inserts. Table and column names
+        are treated as unquoted identifiers, so they are case-insensitive and stored in uppercase,
+        consistent with ``load_table()``.
+
+        Args:
+            df: The data to write. Column names must be valid unquoted identifiers and, unless the
+                table is created by this call, must match the table's columns.
+            table: Table name, optionally qualified, e.g. ``"SCHEMA.TABLE"`` or
+                ``"DB.SCHEMA.TABLE"``. Omitted parts default to the session's current database and
+                schema.
+            overwrite: If ``True``, the existing rows are replaced by ``df``. If ``False``, ``df``
+                is appended. Combined with ``auto_create_table``, the table itself is recreated
+                with column types inferred from ``df``.
+            auto_create_table: If ``True``, the table is created if it does not exist, with column
+                types inferred from ``df``.
+            chunk_size: Number of rows per uploaded file. If ``None``, the whole DataFrame is
+                uploaded at once.
+
+        Returns:
+            The number of rows written. If ``df`` is empty, nothing is written, the table is left
+            unchanged and ``0`` is returned.
+
+        Raises:
+            ValueError: If the table name has more than three parts, or if the table name or a
+                column name is not a valid identifier.
+            RuntimeError: If Snowflake reports that the upload did not complete.
+            SnowflakeError: If the write fails. The error is logged before being re-raised.
+        """
+        parts = [self._check_identifier(p) for p in table.split(".")]
+        if len(parts) > 3:
+            raise ValueError(f"Invalid table name: {table!r}. Expected at most DB.SCHEMA.TABLE.")
+        for col in df.columns:
+            self._check_identifier(str(col))
+
+        if df.empty:
+            logger.warning("DataFrame is empty; nothing written to %s", table)
+            return 0
+
+        # Pad the qualifier on the left, so "T" -> (None, None), "S.T" -> (None, S), etc.
+        *qualifier, table_name = parts
+        database, schema = ([None, None] + qualifier)[-2:]
+
+        conn = self.connect()
+        try:
+            success, n_chunks, n_rows, _ = write_pandas(
+                conn,
+                df,
+                table_name,
+                database=database,
+                schema=schema,
+                chunk_size=chunk_size,
+                quote_identifiers=False,
+                auto_create_table=auto_create_table,
+                overwrite=overwrite,
+                use_logical_type=True,
+            )
+        except SnowflakeError as e:
+            logger.error("Write to %s failed: %s", table, e)
+            raise
+        if not success:
+            raise RuntimeError(f"Write to {table} did not complete successfully.")
+        logger.info("Wrote %d rows to %s in %d chunk(s)", n_rows, table, n_chunks)
+        return n_rows
+
+    def execute(
+        self,
+        sql: str,
+        params: Optional[Union[Sequence[Any], dict[str, Any]]] = None,
+    ) -> Optional[int]:
+        """
+        Runs a SQL statement that does not return data.
+
+        Intended for statements such as ``INSERT``, ``UPDATE``, ``DELETE``, ``MERGE``,
+        ``CREATE TABLE`` or ``USE SCHEMA``. Always pass user-supplied values through ``params``, as
+        in ``query()``. Changes are committed automatically, as the connection runs in Snowflake's
+        default autocommit mode.
+
+        Args:
+            sql: The SQL statement to execute.
+            params: Values bound to the placeholders in ``sql``, as in ``query()``.
+
+        Returns:
+            The number of rows affected, as reported by Snowflake, or ``None`` if unavailable. For
+            statements that do not modify rows (e.g. DDL), the value is not meaningful.
+
+        Raises:
+            SnowflakeError: If the statement fails. The error is logged before being re-raised.
+        """
+        conn = self.connect()
+        try:
+            with conn.cursor() as cur:
+                cur.execute(sql, params)
+                rowcount = cur.rowcount
+        except SnowflakeError as e:
+            logger.error("Statement failed: %s", e)
+            raise
+        logger.info("Statement affected %s rows", rowcount)
+        return rowcount
+
     @staticmethod
     def _check_identifier(name: str) -> str:
         """

{detquantlib-5.4.0 → detquantlib-5.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "detquantlib"
-version = "5.4.0"
+version = "5.5.0"
 description = "An internal library containing functions and classes that can be used across Quant models."
 authors = ["Dynamic Energy Trading"]
 readme = "README.md"