pytrilogy 0.3.142__cp312-cp312-win_amd64.whl

trilogy/dialect/duckdb.py

@@ -0,0 +1,231 @@
+import re
+from typing import Any, Callable, Mapping
+
+from jinja2 import Template
+
+from trilogy.core.enums import FunctionType, Modifier, UnnestMode, WindowType
+from trilogy.core.models.core import DataType
+from trilogy.dialect.base import BaseDialect
+
+WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
+
+SENTINAL_AUTO_CAPTURE_GROUP_VALUE = "-1"
+
+
+def null_wrapper(
+    lval: str,
+    rval: str,
+    modifiers: list[Modifier],
+) -> str:
+
+    if Modifier.NULLABLE in modifiers:
+        return f"{lval} is not distinct from {rval}"
+    return f"{lval} = {rval}"
+
+
+def generate_regex_extract(x: list[str]) -> str:
+    if str(x[2]) == SENTINAL_AUTO_CAPTURE_GROUP_VALUE:
+        regex = re.compile(x[1])
+        if regex.groups == 0:
+            search = 0
+        else:
+            search = 1
+        return f"REGEXP_EXTRACT({x[0]},{x[1]},{search})"
+    return f"REGEXP_EXTRACT({x[0]},{x[1]},{x[2]})"
+
+
+def render_sort(args, types):
+    if len(args) == 1:
+        return f"list_sort({args[0]})"
+    order = args[1].split(" ", 1)
+    if len(order) == 1:
+        return f"list_sort({args[0]}, '{order[0]}')"
+    elif len(order) == 2:
+        return f"list_sort({args[0]}, '{order[0]}', '{order[1]}')"
+
+
+def render_log(args):
+    if len(args) == 1:
+        return f"log({args[0]})"
+    elif len(args) == 2:
+        if int(args[1]) == 10:
+            return f"log({args[0]})"
+        else:
+            # change of base formula
+            return f"log({args[0]})/log({args[1]})"
+    else:
+        raise ValueError("log function requires 1 or 2 arguments")
+
+
+def map_date_part_specifier(specifier: str) -> str:
+    """Map date part specifiers to DuckDB-compatible names"""
+    mapping = {
+        "day_of_week": "dow",
+        # Add other mappings if needed
+    }
+    return mapping.get(specifier, specifier)
+
+
+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.LOG: lambda args, types: render_log(args),
+    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''' {str(a)} ''' for a in args])})"
+    ),
+    FunctionType.SPLIT: lambda args, types: (
+        f"STRING_SPLIT({','.join([f''' {str(a)} ''' for a in args])})"
+    ),
+    ## Duckdb indexes from 1, not 0
+    FunctionType.INDEX_ACCESS: lambda args, types: (f"{args[0]}[{args[1]}]"),
+    ## Duckdb uses list for array
+    FunctionType.ARRAY_DISTINCT: lambda args, types: f"list_distinct({args[0]})",
+    FunctionType.ARRAY_SUM: lambda args, types: f"list_sum({args[0]})",
+    FunctionType.ARRAY_SORT: render_sort,
+    FunctionType.ARRAY_TRANSFORM: lambda args, types: (
+        f"list_transform({args[0]}, {args[1]} -> {args[2]})"
+    ),
+    FunctionType.ARRAY_AGG: lambda args, types: f"array_agg({args[0]})",
+    # datetime is aliased,
+    FunctionType.CURRENT_DATETIME: lambda x, types: "cast(get_current_timestamp() as datetime)",
+    FunctionType.DATETIME: lambda x, types: f"cast({x[0]} as datetime)",
+    FunctionType.TIMESTAMP: lambda x, types: f"cast({x[0]} as timestamp)",
+    FunctionType.DATE: lambda x, types: f"cast({x[0]} as date)",
+    FunctionType.DATE_TRUNCATE: lambda x, types: f"date_trunc('{x[1]}', {x[0]})",
+    FunctionType.DATE_ADD: lambda x, types: f"date_add({x[0]}, {x[2]} * INTERVAL 1 {x[1]})",
+    FunctionType.DATE_SUB: lambda x, types: f"date_add({x[0]}, -{x[2]} * INTERVAL 1 {x[1]})",
+    FunctionType.DATE_PART: lambda x, types: f"date_part('{map_date_part_specifier(x[1])}', {x[0]})",
+    FunctionType.DATE_DIFF: lambda x, types: f"date_diff('{x[2]}', {x[0]}, {x[1]})",
+    FunctionType.CONCAT: lambda x, types: f"({' || '.join(x)})",
+    FunctionType.DATE_LITERAL: lambda x, types: f"date '{x}'",
+    FunctionType.DATETIME_LITERAL: lambda x, types: f"datetime '{x}'",
+    FunctionType.DAY_OF_WEEK: lambda x, types: f"dayofweek({x[0]})",
+    # string
+    FunctionType.CONTAINS: lambda x, types: f"CONTAINS(LOWER({x[0]}), LOWER({x[1]}))",
+    # regexp
+    FunctionType.REGEXP_CONTAINS: lambda x, types: f"REGEXP_MATCHES({x[0]},{x[1]})",
+    FunctionType.REGEXP_EXTRACT: lambda x, types: generate_regex_extract(x),
+}
+
+# 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]}",
+    FunctionType.MAX: lambda args, types: f"{args[0]}",
+    FunctionType.MIN: lambda args, types: f"{args[0]}",
+}
+
+DATATYPE_MAP: dict[DataType, str] = {}
+
+
+DUCKDB_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 %},{% else %}
+{% endif %}{% endfor %}{% endif %}
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}{%- if comment -%}
+-- {{ comment }}
+{%- endif %}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
+    {{ 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 %}
+"""
+)
+
+
+class DuckDBDialect(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,
+    }
+    DATATYPE_MAP = {**BaseDialect.DATATYPE_MAP, **DATATYPE_MAP}
+    QUOTE_CHARACTER = '"'
+    SQL_TEMPLATE = DUCKDB_TEMPLATE
+    UNNEST_MODE = UnnestMode.DIRECT
+    NULL_WRAPPER = staticmethod(null_wrapper)
+
+    def get_table_schema(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[tuple]:
+        """Returns a list of tuples: (column_name, data_type, is_nullable, column_comment)."""
+        column_query = """
+        SELECT
+            column_name,
+            data_type,
+            is_nullable,
+            column_comment
+        FROM information_schema.columns
+        WHERE table_name = ?
+        """
+        params = [table_name]
+
+        if schema:
+            column_query += " AND table_schema = ?"
+            params.append(schema)
+
+        column_query += " ORDER BY ordinal_position"
+
+        # DuckDB supports parameterized queries
+        rows = executor.execute_raw_sql(
+            column_query.replace("?", "'{}'").format(*params)
+        ).fetchall()
+        return rows
+
+    def get_table_primary_keys(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[str]:
+        """Get primary key columns by joining key_column_usage with table_constraints."""
+        pk_query = """
+        SELECT kcu.column_name
+        FROM information_schema.key_column_usage kcu
+        JOIN information_schema.table_constraints tc
+            ON kcu.constraint_name = tc.constraint_name
+            AND kcu.table_name = tc.table_name
+        WHERE kcu.table_name = '{}'
+            AND tc.constraint_type = 'PRIMARY KEY'
+        """.format(
+            table_name
+        )
+
+        if schema:
+            pk_query += " AND kcu.table_schema = '{}'".format(schema)
+
+        pk_query += " ORDER BY kcu.ordinal_position"
+
+        rows = executor.execute_raw_sql(pk_query).fetchall()
+        if rows:
+            return [row[0] for row in rows]
+
+        return []

trilogy/dialect/enums.py

@@ -0,0 +1,147 @@
+from enum import Enum
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+from trilogy.core.models.environment import Environment
+
+if TYPE_CHECKING:
+    from trilogy import Executor
+    from trilogy.hooks.base_hook import BaseHook
+
+from trilogy.constants import Rendering, logger
+from trilogy.dialect.config import DialectConfig, DuckDBConfig
+
+
+def default_factory(conf: DialectConfig, config_type):
+    from sqlalchemy import create_engine
+    from sqlalchemy.pool import NullPool
+
+    # the DuckDB IdentifierPreparer uses a global connection that is not thread safe
+    if isinstance(conf, DuckDBConfig):
+        # we monkey patch to parent to avoid this
+        from duckdb_engine import DuckDBIdentifierPreparer, PGIdentifierPreparer
+
+        DuckDBIdentifierPreparer.__init__ = PGIdentifierPreparer.__init__  # type: ignore
+    engine_args = {
+        "future": True,
+        "poolclass": NullPool,
+    }
+    if not isinstance(conf, config_type):
+        raise TypeError(
+            f"Invalid dialect configuration for type {type(config_type).__name__}, is {type(conf)}"
+        )
+    connect_args = conf.create_connect_args()
+    if connect_args:
+        engine_args["connect_args"] = connect_args
+    return create_engine(conf.connection_string(), **engine_args)
+
+
+class Dialects(Enum):
+    BIGQUERY = "bigquery"
+    SQL_SERVER = "sql_server"
+    DUCK_DB = "duck_db"
+    PRESTO = "presto"
+    TRINO = "trino"
+    POSTGRES = "postgres"
+    SNOWFLAKE = "snowflake"
+    DATAFRAME = "dataframe"
+
+    @classmethod
+    def _missing_(cls, value):
+        if value == "duckdb":
+            return cls.DUCK_DB
+        return super()._missing_(value)
+
+    def default_renderer(self, conf=None, _engine_factory: Callable = default_factory):
+        from trilogy.render import get_dialect_generator
+
+        return get_dialect_generator(self)
+
+    def default_engine(self, conf=None, _engine_factory: Callable = default_factory):
+        if self == Dialects.BIGQUERY:
+            from google.auth import default
+            from google.cloud import bigquery
+
+            from trilogy.dialect.config import BigQueryConfig
+
+            credentials, project = default()
+            client = bigquery.Client(credentials=credentials, project=project)
+            conf = conf or BigQueryConfig(project=project, client=client)
+            return _engine_factory(
+                conf,
+                BigQueryConfig,
+            )
+        elif self == Dialects.SQL_SERVER:
+            raise NotImplementedError()
+        elif self == Dialects.DUCK_DB:
+            from trilogy.dialect.config import DuckDBConfig
+
+            if not conf:
+                conf = DuckDBConfig()
+            return _engine_factory(conf, DuckDBConfig)
+        elif self == Dialects.SNOWFLAKE:
+            from trilogy.dialect.config import SnowflakeConfig
+
+            return _engine_factory(conf, SnowflakeConfig)
+        elif self == Dialects.POSTGRES:
+            logger.warn(
+                "WARN: Using experimental postgres dialect. Most functionality will not work."
+            )
+            import importlib
+
+            spec = importlib.util.find_spec("psycopg2")
+            if spec is None:
+                raise ImportError(
+                    "postgres driver not installed. python -m pip install pypreql[postgres]"
+                )
+            from trilogy.dialect.config import PostgresConfig
+
+            return _engine_factory(conf, PostgresConfig)
+        elif self == Dialects.PRESTO:
+            from trilogy.dialect.config import PrestoConfig
+
+            return _engine_factory(conf, PrestoConfig)
+        elif self == Dialects.TRINO:
+            from trilogy.dialect.config import TrinoConfig
+
+            return _engine_factory(conf, TrinoConfig)
+        elif self == Dialects.DATAFRAME:
+            from trilogy.dialect.config import DataFrameConfig
+            from trilogy.dialect.dataframe import DataframeConnectionWrapper
+
+            if not conf:
+                conf = DataFrameConfig(dataframes={})
+
+            base = _engine_factory(conf, DataFrameConfig)
+
+            return DataframeConnectionWrapper(base, dataframes=conf.dataframes)
+        else:
+            raise ValueError(
+                f"Unsupported dialect {self} for default engine creation; create one explicitly."
+            )
+
+    def default_executor(
+        self,
+        environment: Optional["Environment"] = None,
+        hooks: List["BaseHook"] | None = None,
+        conf: DialectConfig | None = None,
+        rendering: Rendering | None = None,
+        _engine_factory: Callable | None = None,
+    ) -> "Executor":
+        from trilogy import Executor
+
+        if _engine_factory is not None:
+            return Executor(
+                engine=self.default_engine(conf=conf, _engine_factory=_engine_factory),
+                environment=environment or Environment(),
+                dialect=self,
+                rendering=rendering,
+                hooks=hooks,
+            )
+
+        return Executor(
+            engine=self.default_engine(conf=conf),
+            environment=environment or Environment(),
+            dialect=self,
+            rendering=rendering,
+            hooks=hooks,
+        )

trilogy/dialect/metadata.py

@@ -0,0 +1,173 @@
+from typing import Optional
+
+from trilogy.core.enums import DatasourceState, PublishAction
+from trilogy.core.models.datasource import Datasource
+from trilogy.core.models.environment import Environment
+from trilogy.core.statements.author import (
+    ConceptDeclarationStatement,
+    ImportStatement,
+    MergeStatementV2,
+)
+from trilogy.core.statements.execute import (
+    ProcessedPublishStatement,
+    ProcessedShowStatement,
+    ProcessedStaticValueOutput,
+    ProcessedValidateStatement,
+)
+from trilogy.core.validation.common import ValidationTest
+from trilogy.dialect.base import BaseDialect
+from trilogy.dialect.results import MockResult, generate_result_set
+
+
+def handle_concept_declaration(query: ConceptDeclarationStatement) -> MockResult:
+    """Handle concept declaration statements without execution."""
+    concept = query.concept
+    return MockResult(
+        [
+            {
+                "address": concept.address,
+                "type": concept.datatype.value,
+                "purpose": concept.purpose.value,
+                "derivation": concept.derivation.value,
+            }
+        ],
+        ["address", "type", "purpose", "derivation"],
+    )
+
+
+def handle_datasource(query: Datasource) -> MockResult:
+    """Handle datasource queries without execution."""
+    return MockResult(
+        [
+            {
+                "name": query.name,
+            }
+        ],
+        ["name"],
+    )
+
+
+def handle_import_statement(query: ImportStatement) -> MockResult:
+    """Handle import statements without execution."""
+    return MockResult(
+        [
+            {
+                "path": query.path,
+                "alias": query.alias,
+            }
+        ],
+        ["path", "alias"],
+    )
+
+
+def handle_publish_statement(
+    query: ProcessedPublishStatement, environment: Environment
+) -> MockResult:
+    """Handle publish statements by updating environment and returning result."""
+    for x in query.targets:
+        datasource = environment.datasources.get(x)
+        if not datasource:
+            raise ValueError(f"Datasource {x} not found in environment")
+        if query.action == PublishAction.UNPUBLISH:
+            datasource.status = DatasourceState.UNPUBLISHED
+        else:
+            datasource.status = DatasourceState.PUBLISHED
+
+    return MockResult(
+        [{"published": target} for target in query.targets],
+        [
+            "published",
+        ],
+    )
+
+
+def handle_merge_statement(
+    query: MergeStatementV2, environment: Environment
+) -> MockResult:
+    """Handle merge statements by updating environment and returning result."""
+    for concept in query.sources:
+        environment.merge_concept(
+            concept, query.targets[concept.address], modifiers=query.modifiers
+        )
+
+    return MockResult(
+        [
+            {
+                "sources": ",".join([x.address for x in query.sources]),
+                "targets": ",".join([x.address for _, x in query.targets.items()]),
+            }
+        ],
+        ["source", "target"],
+    )
+
+
+def handle_processed_show_statement(
+    query: ProcessedShowStatement, compiled_statements: list[str]
+) -> MockResult:
+    """Handle processed show statements without execution."""
+
+    return generate_result_set(query.output_columns, compiled_statements)
+
+
+def raw_validation_to_result(
+    raw: list[ValidationTest], generator: Optional[BaseDialect] = None
+) -> Optional[MockResult]:
+    """Convert raw validation tests to mock result."""
+    if not raw:
+        return MockResult([], ["check_type", "expected", "result", "ran", "query"])
+    output = []
+    for row in raw:
+        if row.raw_query and generator and not row.generated_query:
+            try:
+                row.generated_query = generator.compile_statement(row.raw_query)
+            except Exception as e:
+                row.generated_query = f"Error generating query: {e}"
+        output.append(
+            {
+                "check_type": row.check_type.value,
+                "expected": row.expected,
+                "result": str(row.result) if row.result else None,
+                "ran": row.ran,
+                "query": row.generated_query if row.generated_query else "",
+            }
+        )
+    return MockResult(output, ["check_type", "expected", "result", "ran", "query"])
+
+
+def handle_processed_validate_statement(
+    query: ProcessedValidateStatement, dialect: BaseDialect, validate_environment_func
+) -> Optional[MockResult]:
+    """Handle processed validate statements."""
+    results = validate_environment_func(query.scope, query.targets)
+    return raw_validation_to_result(results, dialect)
+
+
+def handle_show_statement_outputs(
+    statement: ProcessedShowStatement,
+    compiled_statements: list[str],
+    environment: Environment,
+    dialect: BaseDialect,
+) -> list[MockResult]:
+    """Handle show statement outputs without execution."""
+    output = []
+    for x in statement.output_values:
+        if isinstance(x, ProcessedStaticValueOutput):
+            output.append(generate_result_set(statement.output_columns, x.values))
+        elif compiled_statements:
+
+            output.append(
+                generate_result_set(
+                    statement.output_columns,
+                    compiled_statements,
+                )
+            )
+        elif isinstance(x, ProcessedValidateStatement):
+            from trilogy.core.validation.environment import validate_environment
+
+            raw = validate_environment(environment, x.scope, x.targets)
+            results = raw_validation_to_result(raw, dialect)
+            if results:
+                output.append(results)
+        else:
+            raise NotImplementedError(f"Cannot show type {type(x)} in show statement")
+    return output