pytrilogy 0.3.148__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

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/state_store.py

@@ -0,0 +1,301 @@
+from dataclasses import dataclass, field
+from datetime import date
+
+from sqlalchemy.exc import ProgrammingError
+
+from trilogy import Executor
+from trilogy.core.enums import Purpose
+from trilogy.core.models.build import Factory
+from trilogy.core.models.datasource import (
+    Address,
+    ColumnAssignment,
+    Datasource,
+    RawColumnExpr,
+    UpdateKey,
+    UpdateKeys,
+    UpdateKeyType,
+)
+from trilogy.core.models.environment import Environment
+from trilogy.core.models.execute import CTE
+
+
+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
+
+
+@dataclass
+class DatasourceWatermark:
+    keys: dict[str, UpdateKey]
+
+
+@dataclass
+class StaleAsset:
+    """Represents an asset that needs to be refreshed."""
+
+    datasource_id: str
+    reason: str
+    filters: UpdateKeys = field(default_factory=UpdateKeys)
+
+
+def _compare_watermark_values(
+    a: str | int | float | date, b: str | int | float | date
+) -> int:
+    """Compare two watermark values, returning -1, 0, or 1.
+
+    Handles type mismatches by comparing string representations.
+    """
+    if type(a) is type(b):
+        if a < b:  # type: ignore[operator]
+            return -1
+        elif a > b:  # type: ignore[operator]
+            return 1
+        return 0
+    # Different types: compare as strings
+    sa, sb = str(a), str(b)
+    if sa < sb:
+        return -1
+    elif sa > sb:
+        return 1
+    return 0
+
+
+def get_last_update_time_watermarks(
+    datasource: Datasource, executor: Executor
+) -> DatasourceWatermark:
+    update_time = executor.generator.get_table_last_modified(
+        executor, datasource.safe_address
+    )
+    return DatasourceWatermark(
+        keys={
+            "update_time": UpdateKey(
+                concept_name="update_time",
+                type=UpdateKeyType.UPDATE_TIME,
+                value=update_time,
+            )
+        }
+    )
+
+
+def get_unique_key_hash_watermarks(
+    datasource: Datasource, executor: Executor
+) -> DatasourceWatermark:
+    key_columns: list[ColumnAssignment] = []
+    for col_assignment in datasource.columns:
+        concrete = executor.environment.concepts[col_assignment.concept.address]
+        if concrete.purpose == Purpose.KEY:
+            key_columns.append(col_assignment)
+
+    if not key_columns:
+        return DatasourceWatermark(keys={})
+
+    if isinstance(datasource.address, Address):
+        table_ref = executor.generator.render_source(datasource.address)
+    else:
+        table_ref = datasource.safe_address
+
+    dialect = executor.generator
+    watermarks = {}
+    for col in key_columns:
+        if isinstance(col.alias, str):
+            column_name = col.alias
+        elif isinstance(col.alias, RawColumnExpr):
+            column_name = col.alias.text
+        else:
+            # Function - use rendered expression
+            column_name = str(col.alias)
+        hash_expr = dialect.hash_column_value(column_name)
+        checksum_expr = dialect.aggregate_checksum(hash_expr)
+        query = f"SELECT {checksum_expr} as checksum FROM {table_ref}"
+
+        try:
+            result = executor.execute_raw_sql(query).fetchone()
+            checksum_value = result[0] if result else None
+        except ProgrammingError as e:
+            if _is_table_not_found_error(e, dialect):
+                checksum_value = None
+            else:
+                raise
+
+        watermarks[col.concept.address] = UpdateKey(
+            concept_name=col.concept.address,
+            type=UpdateKeyType.KEY_HASH,
+            value=checksum_value,
+        )
+
+    return DatasourceWatermark(keys=watermarks)
+
+
+def get_incremental_key_watermarks(
+    datasource: Datasource, executor: Executor
+) -> DatasourceWatermark:
+    if not datasource.incremental_by:
+        return DatasourceWatermark(keys={})
+
+    if isinstance(datasource.address, Address):
+        table_ref = executor.generator.render_source(datasource.address)
+    else:
+        table_ref = datasource.safe_address
+
+    watermarks = {}
+    factory = Factory(environment=executor.environment)
+
+    dialect = executor.generator
+    for concept_ref in datasource.incremental_by:
+        concept = executor.environment.concepts[concept_ref.address]
+        build_concept = factory.build(concept)
+        build_datasource = factory.build(datasource)
+        cte: CTE = CTE.from_datasource(build_datasource)
+        # Check if concept is in output_concepts by comparing addresses
+        output_addresses = {c.address for c in datasource.output_concepts}
+        if concept.address in output_addresses:
+            query = f"SELECT MAX({dialect.render_concept_sql(build_concept, cte=cte, alias=False)}) as max_value FROM {table_ref} as {dialect.quote(cte.base_alias)}"
+        else:
+            query = f"SELECT MAX({dialect.render_expr(build_concept.lineage, cte=cte)}) as max_value FROM {table_ref} as {dialect.quote(cte.base_alias)}"
+
+        try:
+            result = executor.execute_raw_sql(query).fetchone()
+            max_value = result[0] if result else None
+        except ProgrammingError as e:
+            if _is_table_not_found_error(e, dialect):
+                max_value = None
+            else:
+                raise
+
+        watermarks[concept.name] = UpdateKey(
+            concept_name=concept.name,
+            type=UpdateKeyType.INCREMENTAL_KEY,
+            value=max_value,
+        )
+
+    return DatasourceWatermark(keys=watermarks)
+
+
+class BaseStateStore:
+
+    def __init__(self) -> None:
+        self.watermarks: dict[str, DatasourceWatermark] = {}
+
+    def watermark_asset(
+        self, datasource: Datasource, executor: Executor
+    ) -> DatasourceWatermark:
+        if datasource.incremental_by:
+            watermarks = get_incremental_key_watermarks(datasource, executor)
+        else:
+            key_columns = [
+                col
+                for col in datasource.columns
+                if executor.environment.concepts[col.concept.address].purpose
+                == Purpose.KEY
+            ]
+            if key_columns:
+                watermarks = get_unique_key_hash_watermarks(datasource, executor)
+            else:
+                watermarks = get_last_update_time_watermarks(datasource, executor)
+
+        self.watermarks[datasource.identifier] = watermarks
+        return watermarks
+
+    def get_datasource_watermarks(
+        self, datasource: Datasource
+    ) -> DatasourceWatermark | None:
+        return self.watermarks.get(datasource.identifier)
+
+    def check_datasource_state(self, datasource: Datasource) -> bool:
+        return datasource.identifier in self.watermarks
+
+    def watermark_all_assets(
+        self, env: Environment, executor: Executor
+    ) -> dict[str, DatasourceWatermark]:
+        """Watermark all datasources in the environment."""
+        for ds in env.datasources.values():
+            self.watermark_asset(ds, executor)
+        return self.watermarks
+
+    def get_stale_assets(
+        self,
+        env: Environment,
+        executor: Executor,
+        root_assets: set[str] | None = None,
+    ) -> list[StaleAsset]:
+        """Find all assets that are stale and need refresh.
+
+        Args:
+            env: The environment containing datasources
+            executor: Executor for querying current state
+            root_assets: Optional set of datasource identifiers that are "source of truth"
+                         and should not be marked stale. If None, uses datasources marked
+                         with is_root=True in the model.
+
+        Returns:
+            List of StaleAsset objects describing what needs refresh and why.
+        """
+        if root_assets is None:
+            root_assets = {
+                ds.identifier for ds in env.datasources.values() if ds.is_root
+            }
+        stale: list[StaleAsset] = []
+
+        # First pass: watermark all assets to get current state
+        self.watermark_all_assets(env, executor)
+
+        # Build map of concept -> max watermark across all assets
+        concept_max_watermarks: dict[str, UpdateKey] = {}
+        for ds_id, watermark in self.watermarks.items():
+            if ds_id in root_assets:
+                # Root assets define the "truth" for incremental keys
+                for key, val in watermark.keys.items():
+                    if (
+                        val.type == UpdateKeyType.INCREMENTAL_KEY
+                        and val.value is not None
+                    ):
+                        existing = concept_max_watermarks.get(key)
+                        if existing is None or (
+                            existing.value is not None
+                            and _compare_watermark_values(val.value, existing.value) > 0
+                        ):
+                            concept_max_watermarks[key] = val
+
+        # Second pass: check non-root assets against max watermarks
+        for ds_id, watermark in self.watermarks.items():
+            if ds_id in root_assets:
+                continue
+
+            for key, val in watermark.keys.items():
+                if val.type == UpdateKeyType.INCREMENTAL_KEY:
+                    max_val = concept_max_watermarks.get(key)
+                    if max_val and max_val.value is not None:
+                        if (
+                            val.value is None
+                            or _compare_watermark_values(val.value, max_val.value) < 0
+                        ):
+                            # Create UpdateKeys with the filter for incremental update
+                            filters = (
+                                UpdateKeys(keys={key: val})
+                                if val.value
+                                else UpdateKeys()
+                            )
+                            stale.append(
+                                StaleAsset(
+                                    datasource_id=ds_id,
+                                    reason=f"incremental key '{key}' behind: {val.value} < {max_val.value}",
+                                    filters=filters,
+                                )
+                            )
+                            break
+
+                elif val.type == UpdateKeyType.UPDATE_TIME:
+                    # For update_time, we'd need root asset update times to compare
+                    # This is tricky without explicit dependency tracking
+                    pass
+
+                elif val.type == UpdateKeyType.KEY_HASH:
+                    # Hash changes indicate data changed, but we need a reference
+                    # to compare against - requires dependency graph
+                    pass
+
+        return stale