pytrilogy 0.3.149__cp313-cp313-win_amd64.whl

trilogy/dialect/snowflake.py

@@ -0,0 +1,129 @@
+from typing import Any, Callable, Mapping
+
+from jinja2 import Template
+
+from trilogy.core.enums import FunctionType, UnnestMode, WindowType
+from trilogy.dialect.base import BaseDialect
+
+ENV_SNOWFLAKE_PW = "PREQL_SNOWFLAKE_PW"
+ENV_SNOWFLAKE_USER = "PREQL_SNOWFLAKE_USER"
+ENV_SNOWFLAKE_ACCOUNT = "PREQL_SNOWFLAKE_ACCOUNT"
+
+WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
+
+FUNCTION_MAP = {
+    FunctionType.COUNT: lambda x, types: f"count({x[0]})",
+    FunctionType.SUM: lambda x, types: f"sum({x[0]})",
+    FunctionType.LENGTH: lambda x, types: f"length({x[0]})",
+    FunctionType.AVG: lambda x, types: f"avg({x[0]})",
+    FunctionType.LIKE: lambda x, types: (
+        f" CASE WHEN {x[0]} like {x[1]} THEN True ELSE False END"
+    ),
+    FunctionType.MINUTE: lambda x, types: f"EXTRACT(MINUTE from {x[0]})",
+    FunctionType.SECOND: lambda x, types: f"EXTRACT(SECOND from {x[0]})",
+    FunctionType.HOUR: lambda x, types: f"EXTRACT(HOUR from {x[0]})",
+    FunctionType.DAY_OF_WEEK: lambda x, types: f"EXTRACT(DAYOFWEEK from {x[0]})",
+    FunctionType.DAY: lambda x, types: f"EXTRACT(DAY from {x[0]})",
+    FunctionType.YEAR: lambda x, types: f"EXTRACT(YEAR from {x[0]})",
+    FunctionType.MONTH: lambda x, types: f"EXTRACT(MONTH from {x[0]})",
+    FunctionType.WEEK: lambda x, types: f"EXTRACT(WEEK from {x[0]})",
+    FunctionType.QUARTER: lambda x, types: f"EXTRACT(QUARTER from {x[0]})",
+    # math
+    FunctionType.POWER: lambda x, types: f"POWER({x[0]}, {x[1]})",
+    FunctionType.DIVIDE: lambda x, types: f"DIV0({x[0]},{x[1]})",
+    FunctionType.UNNEST: lambda x, types: f"table(flatten({x[0]}))",
+    FunctionType.ARRAY: lambda x, types: f"ARRAY_CONSTRUCT({', '.join(x)})",
+}
+
+FUNCTION_GRAIN_MATCH_MAP = {
+    **FUNCTION_MAP,
+    FunctionType.COUNT_DISTINCT: lambda args, types: f"CASE WHEN{args[0]} IS NOT NULL THEN 1 ELSE 0 END",
+    FunctionType.COUNT: lambda args, types: f"CASE WHEN {args[0]} IS NOT NULL THEN 1 ELSE 0 END",
+    FunctionType.SUM: lambda args, types: f"{args[0]}",
+    FunctionType.AVG: lambda args, types: f"{args[0]}",
+}
+
+
+SNOWFLAKE_SQL_TEMPLATE = Template(
+    """{%- if output %}
+{{output}}
+{% endif %}{%- if ctes %}
+WITH {% if recursive%}RECURSIVE{% endif %}{% for cte in ctes %}
+"{{cte.name}}" as ({{cte.statement}}){% if not loop.last %},{% endif %}{% else %}
+{% endfor %}{% endif %}
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
+
+SELECT
+{%- for select in select_columns %}
+    {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
+{% if base %}FROM
+    {{ base }}{% endif %}{% if joins %}{% for join in joins %}
+    {{ join }}{% endfor %}{% endif %}
+{% if where %}WHERE
+    {{ where }}
+{% endif %}
+{%- if group_by %}GROUP BY {% for group in group_by %}
+    {{group}}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}{% if having %}
+HAVING
+\t{{ having }}{% endif %}
+{%- if order_by %}
+ORDER BY {% for order in order_by %}
+    {{ order }}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}
+{%- if limit is not none %}
+LIMIT {{ limit }}{% endif %}{% endif %}
+"""
+)
+MAX_IDENTIFIER_LENGTH = 50
+
+
+class SnowflakeDialect(BaseDialect):
+    WINDOW_FUNCTION_MAP = {**BaseDialect.WINDOW_FUNCTION_MAP, **WINDOW_FUNCTION_MAP}
+    FUNCTION_MAP = {**BaseDialect.FUNCTION_MAP, **FUNCTION_MAP}
+    FUNCTION_GRAIN_MATCH_MAP = {
+        **BaseDialect.FUNCTION_GRAIN_MATCH_MAP,
+        **FUNCTION_GRAIN_MATCH_MAP,
+    }
+    QUOTE_CHARACTER = '"'
+    SQL_TEMPLATE = SNOWFLAKE_SQL_TEMPLATE
+    UNNEST_MODE = UnnestMode.SNOWFLAKE
+
+    def get_table_schema(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[tuple]:
+        """Snowflake requires uppercase identifiers unless quoted."""
+        table_name_upper = table_name.upper()
+
+        column_query = f"""
+        SELECT
+            column_name,
+            data_type,
+            is_nullable,
+            comment as column_comment
+        FROM information_schema.columns
+        WHERE table_name = '{table_name_upper}'
+        """
+        if schema:
+            schema_upper = schema.upper()
+            column_query += f" AND table_schema = '{schema_upper}'"
+        column_query += " ORDER BY ordinal_position"
+
+        rows = executor.execute_raw_sql(column_query).fetchall()
+        return rows
+
+    def get_table_primary_keys(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[str]:
+        """Uses SHOW PRIMARY KEYS; note Snowflake PKs are not enforced."""
+        table_name_upper = table_name.upper()
+
+        # Use SHOW PRIMARY KEYS command (column_name is at index 4)
+        if schema:
+            schema_upper = schema.upper()
+            pk_query = f"SHOW PRIMARY KEYS IN {schema_upper}.{table_name_upper}"
+        else:
+            pk_query = f"SHOW PRIMARY KEYS IN {table_name_upper}"
+
+        rows = executor.execute_raw_sql(pk_query).fetchall()
+        return [row[4] for row in rows] if rows else []

trilogy/dialect/sql_server.py

@@ -0,0 +1,137 @@
+from typing import Any, Callable, Mapping
+
+from jinja2 import Template
+
+from trilogy.core.enums import FunctionType, WindowType
+from trilogy.core.statements.execute import (
+    PROCESSED_STATEMENT_TYPES,
+    ProcessedQuery,
+    ProcessedQueryPersist,
+)
+from trilogy.dialect.base import BaseDialect
+from trilogy.utility import string_to_hash
+
+WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
+
+FUNCTION_MAP = {
+    FunctionType.COUNT: lambda args, types: f"count({args[0]})",
+    FunctionType.SUM: lambda args, types: f"sum({args[0]})",
+    FunctionType.AVG: lambda args, types: f"avg({args[0]})",
+    FunctionType.LENGTH: lambda args, types: f"length({args[0]})",
+    FunctionType.LIKE: lambda args, types: (
+        f" CASE WHEN {args[0]} like {args[1]} THEN True ELSE False END"
+    ),
+    FunctionType.CONCAT: lambda args, types: (
+        f"CONCAT({','.join([f''' '{a}' ''' for a in args])})"
+    ),
+}
+
+# if an aggregate function is called on a source that is at the same grain as the aggregate
+# we may return a static value
+FUNCTION_GRAIN_MATCH_MAP = {
+    **FUNCTION_MAP,
+    FunctionType.COUNT_DISTINCT: lambda args, types: f"CASE WHEN{args[0]} IS NOT NULL THEN 1 ELSE 0 END",
+    FunctionType.COUNT: lambda args, types: f"CASE WHEN {args[0]} IS NOT NULL THEN 1 ELSE 0 END",
+    FunctionType.SUM: lambda args, types: f"{args[0]}",
+    FunctionType.AVG: lambda args, types: f"{args[0]}",
+}
+
+TSQL_TEMPLATE = Template(
+    """{%- if ctes %}
+WITH {% for cte in ctes %}
+{{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
+{%- if full_select -%}{{full_select}}
+{%- else -%}{%- if comment %}
+-- {{ comment }}{%- endif -%}
+SELECT
+{%- if limit is not none %}
+TOP {{ limit }}{% endif %}
+{%- for select in select_columns %}
+    {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
+{% if base %}FROM
+    {{ base }}{% endif %}{% if joins %}
+{% for join in joins %}
+{{ join }}
+{% endfor %}{% endif %}
+{% if where %}WHERE
+    {{ where }}
+{% endif %}
+{%- if group_by %}
+GROUP BY {% for group in group_by %}
+    {{group}}{% if not loop.last %},{% endif %}
+{% endfor %}{% endif %}{% if having %}
+HAVING
+\t{{ having }}{% endif %}
+{%- if order_by %}
+ORDER BY {% for order in order_by %}
+    {{ order }}{% if not loop.last %},{% endif %}
+{% endfor %}{% endif %}{% endif %}
+"""
+)
+
+MAX_IDENTIFIER_LENGTH = 128
+
+
+class SqlServerDialect(BaseDialect):
+    WINDOW_FUNCTION_MAP = {**BaseDialect.WINDOW_FUNCTION_MAP, **WINDOW_FUNCTION_MAP}
+    FUNCTION_MAP = {**BaseDialect.FUNCTION_MAP, **FUNCTION_MAP}
+    FUNCTION_GRAIN_MATCH_MAP = {
+        **BaseDialect.FUNCTION_GRAIN_MATCH_MAP,
+        **FUNCTION_GRAIN_MATCH_MAP,
+    }
+    QUOTE_CHARACTER = '"'
+    SQL_TEMPLATE = TSQL_TEMPLATE
+
+    def get_table_schema(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[tuple]:
+        """Defaults to 'dbo' schema if none specified."""
+        if not schema:
+            schema = "dbo"
+
+        column_query = f"""
+        SELECT
+            column_name,
+            data_type,
+            is_nullable,
+            '' as column_comment
+        FROM information_schema.columns
+        WHERE table_name = '{table_name}'
+        AND table_schema = '{schema}'
+        ORDER BY ordinal_position
+        """
+
+        rows = executor.execute_raw_sql(column_query).fetchall()
+        return rows
+
+    def get_table_primary_keys(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[str]:
+        """Uses sys catalog views for more reliable constraint information."""
+        if not schema:
+            schema = "dbo"
+
+        pk_query = f"""
+        SELECT c.name
+        FROM sys.indexes i
+        INNER JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
+        INNER JOIN sys.columns c ON ic.object_id = c.object_id AND ic.column_id = c.column_id
+        INNER JOIN sys.tables t ON i.object_id = t.object_id
+        INNER JOIN sys.schemas s ON t.schema_id = s.schema_id
+        WHERE i.is_primary_key = 1
+        AND t.name = '{table_name}'
+        AND s.name = '{schema}'
+        ORDER BY ic.key_ordinal
+        """
+
+        rows = executor.execute_raw_sql(pk_query).fetchall()
+        return [row[0] for row in rows]
+
+    def compile_statement(self, query: PROCESSED_STATEMENT_TYPES) -> str:
+        base = super().compile_statement(query)
+        if isinstance(query, (ProcessedQuery, ProcessedQueryPersist)):
+            for cte in query.ctes:
+                if len(cte.name) > MAX_IDENTIFIER_LENGTH:
+                    new_name = f"rhash_{string_to_hash(cte.name)}"
+                    base = base.replace(cte.name, new_name)
+        return base

trilogy/engine.py

@@ -0,0 +1,48 @@
+from typing import Any, Generator, List, Optional, Protocol
+
+from trilogy.core.models.environment import Environment
+
+
+class ResultProtocol(Protocol):
+
+    def fetchall(self) -> List[Any]: ...
+
+    def keys(self) -> List[str]: ...
+
+    def fetchone(self) -> Optional[Any]: ...
+
+    def fetchmany(self, size: int) -> List[Any]: ...
+
+    def __iter__(self) -> Generator[Any, None, None]: ...
+
+
+class EngineConnection(Protocol):
+    pass
+
+    def execute(self, statement: str, parameters: Any | None = None) -> ResultProtocol:
+        pass
+
+    def commit(self):
+        raise NotImplementedError()
+
+    def begin(self):
+        raise NotImplementedError()
+
+    def rollback(self):
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        return
+
+
+class ExecutionEngine(Protocol):
+    pass
+
+    def connect(self) -> EngineConnection:
+        pass
+
+    def setup(self, env: Environment, connection):
+        pass
+
+    def dispose(self, close: bool = True):
+        pass

trilogy/execution/__init__.py

@@ -0,0 +1,17 @@
+from trilogy.dialect.config import (
+    BigQueryConfig,
+    DialectConfig,
+    DuckDBConfig,
+    PostgresConfig,
+    PrestoConfig,
+    SnowflakeConfig,
+)
+
+__all__ = [
+    "DialectConfig",
+    "DuckDBConfig",
+    "PrestoConfig",
+    "SnowflakeConfig",
+    "PostgresConfig",
+    "BigQueryConfig",
+]

trilogy/execution/config.py

@@ -0,0 +1,119 @@
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from tomllib import loads
+
+from trilogy.dialect import (
+    BigQueryConfig,
+    DialectConfig,
+    DuckDBConfig,
+    PostgresConfig,
+    PrestoConfig,
+    SnowflakeConfig,
+    SQLServerConfig,
+)
+from trilogy.dialect.enums import Dialects
+
+DEFAULT_PARALLELISM = 4
+
+
+def load_env_file(env_file_path: Path) -> dict[str, str]:
+    """Load environment variables from a .env file."""
+    env_vars: dict[str, str] = {}
+    if not env_file_path.exists():
+        raise FileNotFoundError(f"Environment file not found: {env_file_path}")
+
+    with open(env_file_path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            if "=" not in line:
+                continue
+            key, _, value = line.partition("=")
+            key = key.strip()
+            value = value.strip()
+            # Remove surrounding quotes if present
+            if (value.startswith('"') and value.endswith('"')) or (
+                value.startswith("'") and value.endswith("'")
+            ):
+                value = value[1:-1]
+            env_vars[key] = value
+    return env_vars
+
+
+def apply_env_vars(env_vars: dict[str, str]) -> None:
+    """Apply environment variables to os.environ."""
+    for key, value in env_vars.items():
+        os.environ[key] = value
+
+
+@dataclass
+class RuntimeConfig:
+
+    startup_trilogy: list[Path]
+    startup_sql: list[Path]
+    parallelism: int = DEFAULT_PARALLELISM
+    engine_dialect: Dialects | None = None
+    engine_config: DialectConfig | None = None
+    source_path: Path | None = None
+    env_files: list[Path] = field(default_factory=list)
+
+
+def load_config_file(path: Path) -> RuntimeConfig:
+    with open(path, "r") as f:
+        toml_content = f.read()
+    config_data = loads(toml_content)
+
+    engine_raw: dict = config_data.get("engine", {})
+    engine_config_raw = engine_raw.get("config", {})
+    engine = Dialects(engine_raw.get("dialect")) if engine_raw.get("dialect") else None
+    engine_config: DialectConfig | None
+    if engine:
+        if engine == Dialects.DUCK_DB:
+            engine_config = (
+                DuckDBConfig(**engine_config_raw) if engine_config_raw else None
+            )
+        elif engine == Dialects.POSTGRES:
+            engine_config = (
+                PostgresConfig(**engine_config_raw) if engine_config_raw else None
+            )
+        elif engine == Dialects.PRESTO:
+            engine_config = (
+                PrestoConfig(**engine_config_raw) if engine_config_raw else None
+            )
+        elif engine == Dialects.SNOWFLAKE:
+            engine_config = (
+                SnowflakeConfig(**engine_config_raw) if engine_config_raw else None
+            )
+        elif engine == Dialects.SQL_SERVER:
+            engine_config = (
+                SQLServerConfig(**engine_config_raw) if engine_config_raw else None
+            )
+        elif engine == Dialects.BIGQUERY:
+            engine_config = (
+                BigQueryConfig(**engine_config_raw) if engine_config_raw else None
+            )
+        else:
+            engine_config = None
+    else:
+        engine_config = None
+    setup: dict = config_data.get("setup", {})
+
+    # Parse env_file - can be a single string or list of strings
+    env_raw = engine_raw.get("env_file", [])
+    if isinstance(env_raw, str):
+        env_files = [path.parent / env_raw]
+    else:
+        env_files = [path.parent / p for p in env_raw]
+
+    return RuntimeConfig(
+        startup_trilogy=[path.parent / p for p in setup.get("trilogy", [])],
+        startup_sql=[path.parent / p for p in setup.get("sql", [])],
+        parallelism=config_data.get("parallelism", DEFAULT_PARALLELISM),
+        engine_dialect=engine,
+        engine_config=engine_config,
+        source_path=path,
+        env_files=env_files,
+    )

trilogy/execution/state/exceptions.py

@@ -0,0 +1,26 @@
+from sqlalchemy.exc import ProgrammingError
+
+
+def _is_table_not_found_error(exc: ProgrammingError, dialect) -> bool:
+    """Check if exception is a table-not-found error for the given dialect."""
+    pattern = dialect.TABLE_NOT_FOUND_PATTERN
+    if pattern is None:
+        return False
+    error_msg = str(exc.orig) if exc.orig else str(exc)
+    return pattern in error_msg
+
+
+def _is_http_not_found_error(exc: Exception, dialect) -> bool:
+    """Check if exception is an HTTP 404 error (e.g., GCS file not found)."""
+    pattern = getattr(dialect, "HTTP_NOT_FOUND_PATTERN", None)
+    if pattern is None:
+        return False
+    error_msg = str(exc.orig) if hasattr(exc, "orig") and exc.orig else str(exc)
+    return pattern in error_msg
+
+
+def is_missing_source_error(exc: Exception, dialect) -> bool:
+    """Check if exception indicates a missing source (table or remote file)."""
+    if isinstance(exc, ProgrammingError):
+        return _is_table_not_found_error(exc, dialect)
+    return _is_http_not_found_error(exc, dialect)