rivage-clair 0.1.1__py3-none-any.whl

clair/__init__.py

@@ -0,0 +1,98 @@
+"""Clair -- Python-native data transformation framework for Snowflake.
+
+Public API exports. Users import from here in their Trouve files:
+
+    from clair import Trouve, TrouveType, Column, ColumnType
+
+Runtime context
+---------------
+``clair.env`` is set to the active :class:`~clair.environments.environments.Environment`
+before Trouve modules are loaded during discovery. Import it to implement
+feature flags based on the active environment::
+
+    import clair
+
+    trouve = Trouve(
+        sql=f"SELECT * FROM {upstream} {'WHERE is_beta = 1' if clair.env.role == 'DEV' else ''}"
+    )
+
+``clair.run_mode`` is set to the active :class:`~clair.trouves.run_config.RunMode`
+before Trouve modules are loaded during discovery. Use it to make SQL conditional
+on run mode (analogous to dbt's ``is_incremental()``)::
+
+    import clair
+    from clair import RunMode
+
+    trouve = Trouve(
+        sql=f\"\"\"
+            SELECT * FROM {upstream}
+            {'WHERE created_at > dateadd(\\'day\\', -3, current_timestamp())' if clair.run_mode == RunMode.INCREMENTAL else ''}
+        \"\"\"
+    )
+
+When ``clair.run_mode`` is ``None`` (e.g. during ``clair dag`` or ``clair docs``),
+the expression ``clair.run_mode == RunMode.INCREMENTAL`` evaluates to ``False``,
+so no WHERE clause is applied — the safe default for non-run commands.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from clair.environments.environments import Environment
+
+from clair.trouves.column import Column, ColumnType
+from clair.trouves.config import DatabaseDefaults, SchemaDefaults
+from clair.trouves.run_config import (
+    IncrementalMode,
+    UpsertConfig,
+    RunConfig,
+    RunMode,
+    SOURCE,
+    TARGET,
+)
+from clair.trouves.test import (
+    AnyTest,
+    Test,
+    TestNotNull,
+    TestRowCount,
+    TestSql,
+    TestUnique,
+    TestUniqueColumns,
+    THIS,
+)
+from clair.trouves.trouve import Trouve, TrouveType
+
+__version__ = "0.1.0"
+
+# Set by discover_project() before Trouve modules are loaded.
+# None when running outside of a clair discovery run.
+env: Environment | None = None
+
+# Set by discover_project() before Trouve modules are loaded.
+# None when running outside of a clair discovery run.
+run_mode: RunMode | None = None
+
+__all__ = [
+    "AnyTest",
+    "Column",
+    "ColumnType",
+    "DatabaseDefaults",
+    "IncrementalMode",
+    "UpsertConfig",
+    "RunConfig",
+    "RunMode",
+    "SOURCE",
+    "SchemaDefaults",
+    "TARGET",
+    "Test",
+    "TestNotNull",
+    "TestRowCount",
+    "TestSql",
+    "TestUnique",
+    "TestUniqueColumns",
+    "THIS",
+    "Trouve",
+    "TrouveType",
+]

clair/_logging.py

@@ -0,0 +1,61 @@
+"""Structlog configuration for the Clair CLI."""
+
+from __future__ import annotations
+
+import os
+import sys
+from collections.abc import MutableMapping
+from datetime import datetime
+from typing import Any
+
+import click
+import structlog
+
+_LEVEL_FG = {"debug": "blue", "info": "green", "warning": "yellow", "error": "red", "critical": "red"}
+
+
+def _multiline_renderer(_logger: Any, _method: str, event_dict: MutableMapping[str, Any]) -> str:
+    level = event_dict.pop("level", "info")
+    event = event_dict.pop("event", "")
+    event_dict.pop("timestamp", None)  # replaced by our own below
+
+    ts = datetime.now().strftime("%H:%M:%S.%f")[:-3]  # millisecond precision
+    kv = {"timestamp": ts, **event_dict}
+
+    if sys.stderr.isatty():
+        level_str = click.style(f"{level:<8}", fg=_LEVEL_FG.get(level, "white"), bold=True)
+        header = f"[{level_str}] {click.style(event, bold=True)}"
+        kv_lines = "\n".join(f"  {click.style(k, fg='cyan')}={click.style(str(v), fg='magenta')}" for k, v in kv.items())
+    else:
+        header = f"[{level:<8}] {event}"
+        kv_lines = "\n".join(f"  {k}={v}" for k, v in kv.items())
+
+    return f"{header}\n{kv_lines}" if kv_lines else header
+
+
+def configure_logging() -> None:
+    log_format = os.environ.get("CLAIR_LOG_FORMAT", "").lower()
+    use_json = log_format == "json"
+
+    if use_json:
+        structlog.configure(
+            processors=[
+                structlog.stdlib.add_log_level,
+                structlog.processors.TimeStamper(fmt="iso", utc=True),
+                structlog.processors.JSONRenderer(),
+            ],
+            wrapper_class=structlog.make_filtering_bound_logger(20),  # INFO
+            context_class=dict,
+            logger_factory=structlog.PrintLoggerFactory(file=sys.stdout),
+        )
+    else:
+        structlog.configure(
+            processors=[
+                structlog.stdlib.add_log_level,
+                structlog.processors.TimeStamper(fmt="%H:%M:%S", utc=False),
+                _multiline_renderer,
+            ],
+            wrapper_class=structlog.make_filtering_bound_logger(20),  # INFO
+            context_class=dict,
+            logger_factory=structlog.PrintLoggerFactory(file=sys.stderr),
+        )

clair/adapters/base.py

@@ -0,0 +1,82 @@
+"""Abstract warehouse adapter interface.
+
+Abstracts the database connection so that future adapters (e.g. BigQuery,
+Databricks) can be added without rewriting the runner.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import pandas as pd
+from pydantic import BaseModel
+
+
+class QueryResult(BaseModel):
+    """Result of executing a single SQL query against the warehouse.
+
+    Attributes:
+        query_id: Warehouse-assigned identifier for the executed query.
+        query_url: URL to the query detail page in the warehouse console.
+        success: Whether the query completed without error.
+        error: Error message if the query failed.
+        row_count: Number of rows returned (or affected) by the query.
+    """
+
+    query_id: str
+    query_url: str
+    success: bool
+    error: str | None = None
+    row_count: int = 0
+
+
+class WarehouseAdapter(ABC):
+    """Abstract interface for warehouse connections."""
+
+    @abstractmethod
+    def connect(self, profile: dict[str, Any]) -> None:
+        """Establish a connection using profile credentials."""
+        ...
+
+    @abstractmethod
+    def execute(self, sql: str) -> QueryResult:
+        """Execute a SQL statement and return the result."""
+        ...
+
+    @abstractmethod
+    def set_context(
+        self,
+        warehouse: str | None = None,
+        role: str | None = None,
+        database_name: str | None = None,
+    ) -> None:
+        """Set the session context (warehouse, role, database)."""
+        ...
+
+    @abstractmethod
+    def table_exists(self, database_name: str, schema_name: str, table_name: str) -> bool:
+        """Check whether a table exists in the warehouse."""
+        ...
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the connection."""
+        ...
+
+    @abstractmethod
+    def fetch_dataframe(self, full_name: str) -> pd.DataFrame:
+        """Fetch a table as a pandas DataFrame."""
+        ...
+
+    @abstractmethod
+    def write_dataframe(
+        self,
+        dataframe: pd.DataFrame,
+        full_name: str,
+        database_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> QueryResult:
+        """Write a DataFrame to the warehouse, creating or replacing the table."""
+        ...

clair/adapters/snowflake.py

@@ -0,0 +1,192 @@
+"""Snowflake adapter implementation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+import snowflake.connector
+from snowflake.connector.pandas_tools import write_pandas
+from cryptography.hazmat.primitives import serialization
+
+from clair.adapters.base import QueryResult, WarehouseAdapter
+
+
+class SnowflakeAdapter(WarehouseAdapter):
+    """Snowflake warehouse adapter using snowflake-connector-python."""
+
+    def __init__(self) -> None:
+        self._conn: snowflake.connector.SnowflakeConnection | None = None
+        self._region: str = ""
+        self._account_locator: str = ""
+
+    def connect(self, profile: dict[str, Any]) -> None:
+        """Connect to Snowflake using profile credentials.
+
+        Supports:
+        - SSO via authenticator=externalbrowser
+        - Key pair auth via private_key_path
+        - Standard username/password
+        """
+        self._region = profile.get("region", "")
+        self._account_locator = profile.get("account_locator", "")
+
+        connect_args: dict[str, Any] = {
+            "account": profile["account"],
+            "user": profile["user"],
+        }
+
+        # Auth method
+        if "authenticator" in profile:
+            connect_args["authenticator"] = profile["authenticator"]
+        elif "private_key_pem" in profile:
+            pem_content = profile["private_key_pem"]
+            passphrase = profile.get("private_key_passphrase")
+            password = passphrase.encode() if isinstance(passphrase, str) else passphrase
+            pem_bytes = pem_content.encode() if isinstance(pem_content, str) else pem_content
+            p_key = serialization.load_pem_private_key(pem_bytes, password=password)
+            connect_args["private_key"] = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+        elif "private_key_path" in profile:
+            key_path = Path(profile["private_key_path"]).expanduser()
+            passphrase = profile.get("private_key_passphrase")
+            password = passphrase.encode() if isinstance(passphrase, str) else passphrase
+            with open(key_path, "rb") as f:
+                p_key = serialization.load_pem_private_key(f.read(), password=password)
+                connect_args["private_key"] = p_key.private_bytes(
+                    encoding=serialization.Encoding.DER,
+                    format=serialization.PrivateFormat.PKCS8,
+                    encryption_algorithm=serialization.NoEncryption(),
+                )
+        elif "password" in profile:
+            connect_args["password"] = profile["password"]
+
+        # Optional context
+        for key in ("warehouse", "role", "database"):
+            if key in profile:
+                connect_args[key] = profile[key]
+
+        self._conn = snowflake.connector.connect(**connect_args)
+
+    def execute(self, sql: str) -> QueryResult:
+        """Execute SQL and return a QueryResult with query ID and URL."""
+        if self._conn is None:
+            raise RuntimeError("Not connected. Call connect() first.")
+
+        cursor = self._conn.cursor()
+        try:
+            cursor.execute(sql)
+            query_id = cursor.sfqid or "unknown"
+            return QueryResult(
+                query_id=query_id,
+                query_url=self._build_query_url(query_id),
+                success=True,
+                row_count=cursor.rowcount or 0,
+            )
+        except Exception as e:
+            query_id = getattr(cursor, "sfqid", None) or "unknown"
+            return QueryResult(
+                query_id=query_id,
+                query_url=self._build_query_url(query_id),
+                success=False,
+                error=str(e),
+            )
+        finally:
+            cursor.close()
+
+    def table_exists(self, database_name: str, schema_name: str, table_name: str) -> bool:
+        """Check whether a table exists in Snowflake via INFORMATION_SCHEMA."""
+        result = self.execute(
+            f"SELECT 1 FROM {database_name}.INFORMATION_SCHEMA.TABLES "
+            f"WHERE TABLE_CATALOG = '{database_name.upper()}' "
+            f"AND TABLE_SCHEMA = '{schema_name.upper()}' "
+            f"AND TABLE_NAME = '{table_name.upper()}'"
+        )
+        return result.row_count > 0
+
+    def set_context(
+        self,
+        warehouse: str | None = None,
+        role: str | None = None,
+        database_name: str | None = None,
+    ) -> None:
+        """Set session context by executing USE commands.
+
+        Only issues USE statements for values that are non-None and non-empty.
+        Order: ROLE first (affects permissions), then WAREHOUSE, then DATABASE.
+        """
+        if self._conn is None:
+            raise RuntimeError("Not connected. Call connect() first.")
+
+        cursor = self._conn.cursor()
+        try:
+            if role and role.strip():
+                cursor.execute(f"USE ROLE {role}")
+            if warehouse and warehouse.strip():
+                cursor.execute(f"USE WAREHOUSE {warehouse}")
+            if database_name and database_name.strip():
+                cursor.execute(f"USE DATABASE {database_name}")
+        finally:
+            cursor.close()
+
+    def fetch_dataframe(self, full_name: str) -> pd.DataFrame:
+        """Fetch an entire Snowflake table as a pandas DataFrame."""
+        if self._conn is None:
+            raise RuntimeError("Not connected. Call connect() first.")
+
+        cursor = self._conn.cursor()
+        try:
+            cursor.execute(f"SELECT * FROM {full_name}")
+            dataframe = cursor.fetch_pandas_all()
+            dataframe.columns = dataframe.columns.str.lower()
+            return dataframe
+        finally:
+            cursor.close()
+
+    def write_dataframe(
+        self,
+        dataframe: pd.DataFrame,
+        full_name: str,
+        database_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> QueryResult:
+        """Write a DataFrame to Snowflake, creating or replacing the table."""
+        if self._conn is None:
+            raise RuntimeError("Not connected. Call connect() first.")
+
+        success, num_chunks, num_rows, output = write_pandas(
+            conn=self._conn,
+            df=dataframe,
+            table_name=table_name.upper(),
+            database=database_name.upper(),
+            schema=schema_name.upper(),
+            auto_create_table=True,
+            overwrite=True,
+            quote_identifiers=False,
+        )
+        # query_id and query_url are empty: write_dataframe uses CREATE TEMP STAGE
+        # and PUT under the hood, not a single queryable SQL statement.
+        return QueryResult(
+            query_id="",
+            query_url="",
+            success=success,
+            row_count=num_rows,
+        )
+
+    def close(self) -> None:
+        """Close the Snowflake connection."""
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    def _build_query_url(self, query_id: str) -> str:
+        """Build a Snowflake console URL for a query ID."""
+        return (
+            f"https://app.snowflake.com/{self._region}/{self._account_locator}"
+            f"/#/compute/history/queries/{query_id}/detail"
+        )