pytrilogy 0.3.138__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

trilogy/dialect/duckdb.py

@@ -0,0 +1,177 @@
+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)

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

trilogy/dialect/mock.py

@@ -0,0 +1,190 @@
+import random
+from datetime import date, datetime
+from typing import TYPE_CHECKING, Any, Iterable
+
+from trilogy.core.enums import Purpose
+from trilogy.core.models.author import Concept, ConceptRef
+from trilogy.core.models.core import CONCRETE_TYPES, ArrayType, DataType, TraitDataType
+from trilogy.core.models.datasource import Address, Datasource
+from trilogy.core.models.environment import Environment
+from trilogy.core.statements.execute import ProcessedMockStatement
+from trilogy.dialect.results import MockResult
+
+if TYPE_CHECKING:
+    from pyarrow import Table
+
+DEFAULT_SCALE_FACTOR = 100
+
+
+def safe_name(name: str) -> str:
+    return "".join(c if c.isalnum() or c == "_" else "_" for c in name)
+
+
+def mock_email(scale_factor: int, is_key: bool = False) -> list[str]:
+    providers = ["example.com", "test.com", "mock.com", "sample.org"]
+    if is_key:
+        return [
+            f"user{i}@{providers[i % len(providers)]}"
+            for i in range(1, scale_factor + 1)
+        ]
+    return [
+        f"user{random.randint(1, 999999)}@{random.choice(providers)}"
+        for _ in range(scale_factor)
+    ]
+
+
+def mock_hex_code(scale_factor: int, is_key: bool = False) -> list[str]:
+    if is_key:
+        return [f"#{i:06x}" for i in range(1, scale_factor + 1)]
+    return [f"#{random.randint(0, 0xFFFFFF):06x}" for _ in range(scale_factor)]
+
+
+def mock_datatype(
+    full_type: Any, datatype: CONCRETE_TYPES, scale_factor: int, is_key: bool = False
+) -> list[Any]:
+    if isinstance(full_type, TraitDataType):
+        if full_type.type == DataType.STRING:
+            # TODO: get stdlib inventory some other way?
+            if full_type.traits == ["email_address"]:
+                # email mock function
+                return mock_email(scale_factor, is_key)
+            elif full_type.traits == ["hex"]:
+                return mock_hex_code(scale_factor, is_key)
+        return mock_datatype(full_type.type, full_type.type, scale_factor, is_key)
+    elif datatype == DataType.INTEGER:
+        if is_key:
+            # unique integers for keys
+            return list(range(1, scale_factor + 1))
+        return [random.randint(0, 999_999) for _ in range(scale_factor)]
+    elif datatype == DataType.STRING:
+        if is_key:
+            # unique strings for keys
+            return [f"key_{i}" for i in range(1, scale_factor + 1)]
+        return [
+            f"mock_string_{random.randint(0, 999_999)}" for _ in range(scale_factor)
+        ]
+    elif datatype == DataType.FLOAT:
+        if is_key:
+            # unique floats for keys
+            return [float(i) for i in range(1, scale_factor + 1)]
+        return [random.uniform(0, 999_999) for _ in range(scale_factor)]
+    elif datatype == DataType.NUMERIC:
+        if is_key:
+            # unique numerics for keys
+            return [float(i) for i in range(1, scale_factor + 1)]
+        return [round(random.uniform(0, 999_999), 2) for _ in range(scale_factor)]
+    elif datatype == DataType.BOOL:
+        # booleans can only have 2 unique values, so keys don't make sense here
+        return [random.choice([True, False]) for _ in range(scale_factor)]
+    elif datatype == DataType.DATE:
+        if is_key:
+            # unique dates for keys - spread across multiple months/years if needed
+            base_date = date(2023, 1, 1)
+            return [
+                date.fromordinal(base_date.toordinal() + i) for i in range(scale_factor)
+            ]
+        return [date(2023, 1, random.randint(1, 28)) for _ in range(scale_factor)]
+    elif datatype in (DataType.DATETIME, DataType.TIMESTAMP):
+        if is_key:
+            # unique datetimes for keys - increment by seconds
+            base_dt = datetime(2023, 1, 1, 0, 0, 0)
+            return [
+                datetime.fromtimestamp(base_dt.timestamp() + i)
+                for i in range(scale_factor)
+            ]
+        return [
+            datetime(
+                2023,
+                1,
+                1,
+                random.randint(0, 23),
+                random.randint(0, 59),
+                random.randint(0, 59),
+            )
+            for _ in range(scale_factor)
+        ]
+    elif isinstance(datatype, ArrayType):
+        # arrays as keys don't typically make sense, but generate unique if requested
+        if is_key:
+            return [
+                [mock_datatype(datatype.type, datatype.value_data_type, 5, False)[0], i]
+                for i in range(scale_factor)
+            ]
+        return [
+            [mock_datatype(datatype.type, datatype.value_data_type, 5, False)]
+            for _ in range(scale_factor)
+        ]
+    raise NotImplementedError(f"Mocking not implemented for datatype {datatype}")
+
+
+class MockManager:
+
+    def __init__(
+        self, environment: Environment, scale_factor: int = DEFAULT_SCALE_FACTOR
+    ):
+        self.environment = environment
+        self.concept_mocks: dict[str, Any] = {}
+        self.scale_factor = scale_factor
+
+    def mock_concept(self, concept: Concept | ConceptRef):
+        if concept.address in self.concept_mocks:
+            return False
+        concrete = self.environment.concepts[concept.address]
+        self.concept_mocks[concept.address] = mock_datatype(
+            concept.datatype,
+            concept.output_datatype,
+            self.scale_factor,
+            True if concrete.purpose == Purpose.KEY else False,
+        )
+        return True
+
+    def create_mock_table(
+        self, concepts: Iterable[Concept | ConceptRef], headers: list[str]
+    ) -> "Table":
+        from pyarrow import table
+
+        data = {h: self.concept_mocks[c.address] for h, c in zip(headers, concepts)}
+        return table(data)
+
+
+def handle_processed_mock_statement(
+    query: ProcessedMockStatement, environment: Environment, executor
+) -> MockResult:
+    """Handle processed mock statements."""
+    # For mock statements, we can simulate some output based on targets
+    mock_manager = MockManager(environment)
+    output = []
+    for target in query.targets:
+        datasource = environment.datasources.get(target)
+        if not datasource:
+            raise ValueError(f"Datasource {target} not found in environment")
+        mock_datasource(datasource, mock_manager, executor)
+        output.append(
+            {
+                "target": target,
+                "status": "mocked",
+            }
+        )
+    return MockResult(output, ["target", "status"])
+
+
+def mock_datasource(datasource: Datasource, manager: MockManager, executor):
+    concrete: list[ConceptRef] = []
+    headers: list[str] = []
+    for k, col in datasource.concrete_columns.items():
+        manager.mock_concept(col.concept)
+        concrete.append(col.concept)
+        headers.append(k)
+
+    table = manager.create_mock_table(concrete, headers)
+
+    # duckdb load the pyarrow table
+    executor.execute_raw_sql(
+        "register(:name, :tbl)", {"name": "mock_tbl", "tbl": table}
+    )
+    address = safe_name(datasource.safe_address)
+    executor.execute_raw_sql(
+        f"""CREATE OR REPLACE TABLE {address} AS SELECT * FROM mock_tbl"""
+    )
+    # overwrite the address since we've mangled the name
+    datasource.address = Address(location=address)