pytrilogy 0.3.142__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

trilogy/dialect/bigquery.py

@@ -0,0 +1,308 @@
+import uuid
+from typing import Any, Callable, Dict, Mapping, Optional
+
+from jinja2 import Template
+
+from trilogy.core.enums import (
+    ComparisonOperator,
+    FunctionType,
+    UnnestMode,
+    WindowType,
+)
+from trilogy.core.models.core import (
+    DataType,
+)
+from trilogy.core.models.execute import CTE, CompiledCTE, UnionCTE
+from trilogy.core.statements.execute import ProcessedQueryPersist
+from trilogy.dialect.base import BaseDialect, safe_quote
+
+WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
+
+
+def transform_date_part(part: str) -> str:
+    part_upper = part.upper()
+    if part_upper == "DAY_OF_WEEK":
+        return "DAYOFWEEK"
+    return part_upper
+
+
+def handle_length(args, types: list[DataType] | None = None) -> str:
+    arg = args[0]
+    if types and types[0].data_type == DataType.ARRAY:
+        return f"ARRAY_LENGTH({arg})"
+    return f"LENGTH({arg})"
+
+
+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: handle_length(x, types),
+    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.IS_NULL: lambda x, types: f"{x[0]} IS NULL",
+    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]})-1",  # BigQuery's DAYOFWEEK returns 1 for Sunday
+    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"COALESCE(SAFE_DIVIDE({x[0]},{x[1]}),0)",
+    FunctionType.DATE_ADD: lambda x, types: f"DATE_ADD({x[0]}, INTERVAL {x[2]} {x[1]})",
+    FunctionType.DATE_SUB: lambda x, types: f"DATE_SUB({x[0]}, INTERVAL {x[2]} {x[1]})",
+    FunctionType.DATE_PART: lambda x, types: f"EXTRACT({transform_date_part(x[1])} FROM {x[0]})",
+    FunctionType.MONTH_NAME: lambda x, types: f"FORMAT_DATE('%B', {x[0]})",
+    FunctionType.DAY_NAME: lambda x, types: f"FORMAT_DATE('%A', {x[0]})",
+    # string
+    FunctionType.CONTAINS: lambda x, types: f"CONTAINS_SUBSTR({x[0]}, {x[1]})",
+    FunctionType.RANDOM: lambda x, types: f"FLOOR(RAND()*{x[0]})",
+    FunctionType.ARRAY_SUM: lambda x, types: f"(select sum(x) from unnest({x[0]}) as x)",
+    FunctionType.ARRAY_DISTINCT: lambda x, types: f"ARRAY(SELECT DISTINCT element FROM UNNEST({x[0]}) AS element)",
+    FunctionType.ARRAY_SORT: lambda x, types: f"ARRAY(SELECT element FROM UNNEST({x[0]}) AS element ORDER BY element)",
+    # aggregate
+    FunctionType.BOOL_AND: lambda x, types: f"LOGICAL_AND({x[0]})",
+    FunctionType.BOOL_OR: lambda x, types: f"LOGICAL_OR({x[0]})",
+}
+
+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]}",
+}
+
+DATATYPE_MAP: dict[DataType, str] = {
+    DataType.STRING: "STRING",
+    DataType.INTEGER: "INT64",
+    DataType.FLOAT: "FLOAT64",
+    DataType.BOOL: "BOOL",
+    DataType.NUMERIC: "NUMERIC",
+    DataType.MAP: "MAP",
+    DataType.DATE: "DATE",
+    DataType.DATETIME: "DATETIME",
+    DataType.TIMESTAMP: "TIMESTAMP",
+}
+
+
+BQ_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 %},{% else%}
+{% endif %}{% 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 %}
+"""
+)
+
+
+BQ_CREATE_TABLE_SQL_TEMPLATE = Template(
+    """
+CREATE {% if create_mode == "create_or_replace" %}OR REPLACE TABLE{% elif create_mode == "create_if_not_exists" %}TABLE IF NOT EXISTS{% else %}TABLE{% endif %} {{ name}} (
+{%- for column in columns %}
+    `{{ column.name }}` {{ type_map[column.name] }}{% if column.description %} OPTIONS(description='{{ column.description }}'){% endif %}{% if not loop.last %},{% endif %}
+{%- endfor %}
+)
+{%- if partition_by %}
+PARTITION BY {{ partition_by }}
+{%- endif %}
+{%- if cluster_by %}
+CLUSTER BY {{ cluster_by | join(', ') }}
+{%- endif %}
+{%- if table_description %}
+OPTIONS(
+    description='{{ table_description }}'
+)
+{%- endif %};
+""".strip()
+)
+
+PARTITIONED_INSERT_TEMPLATE = Template(
+    """
+-- Step 1: materialize results
+CREATE TEMP TABLE {{ tmp_table }} AS SELECT * FROM  {{ target_table }} limit 0;
+                                       
+INSERT INTO {{ tmp_table }}
+    {{ final_select }}
+;
+
+-- Step 2: extract distinct partitions and generate dynamic statements
+BEGIN
+    DECLARE partition_values ARRAY<{{ partition_type }}>;
+    DECLARE current_partition {{ partition_type }};
+    DECLARE i INT64 DEFAULT 0;
+    
+    -- Get all distinct partition values
+    SET partition_values = (
+        SELECT ARRAY_AGG(DISTINCT {{ partition_key[0] }})
+        FROM {{ tmp_table }}
+    );
+    
+    -- Loop through each partition value
+    WHILE i < ARRAY_LENGTH(partition_values) DO
+        SET current_partition = partition_values[OFFSET(i)];
+        
+        -- Delete existing records for this partition
+        EXECUTE IMMEDIATE FORMAT(
+            'DELETE FROM {{ target_table }} WHERE {{ partition_key[0] }} = "%t"',
+            current_partition
+        );
+        
+        -- Insert new records for this partition
+        EXECUTE IMMEDIATE FORMAT(
+            'INSERT INTO {{ target_table }} SELECT * FROM {{ tmp_table }} WHERE {{ partition_key[0] }} = "%t"',
+            current_partition
+        );
+        
+        SET i = i + 1;
+    END WHILE;
+END;
+"""
+)
+
+MAX_IDENTIFIER_LENGTH = 50
+
+
+def parse_bigquery_table_name(
+    table_name: str, schema: str | None = None
+) -> tuple[str, str | None]:
+    """Parse BigQuery table names supporting project.dataset.table format."""
+    if "." in table_name and not schema:
+        parts = table_name.split(".")
+        if len(parts) == 2:
+            schema = parts[0]
+            table_name = parts[1]
+        elif len(parts) == 3:
+            # project.dataset.table format
+            schema = f"{parts[0]}.{parts[1]}"
+            table_name = parts[2]
+    return table_name, schema
+
+
+class BigqueryDialect(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 = BQ_SQL_TEMPLATE
+    CREATE_TABLE_SQL_TEMPLATE = BQ_CREATE_TABLE_SQL_TEMPLATE
+    UNNEST_MODE = UnnestMode.CROSS_JOIN_UNNEST
+    DATATYPE_MAP = DATATYPE_MAP
+
+    def get_table_schema(
+        self, executor, table_name: str, schema: str | None = None
+    ) -> list[tuple]:
+        """BigQuery uses dataset instead of schema and supports project.dataset.table format."""
+        table_name, schema = parse_bigquery_table_name(table_name, schema)
+
+        column_query = f"""
+        SELECT
+            column_name,
+            data_type,
+            is_nullable,
+            '' as column_comment
+        FROM `{schema}.INFORMATION_SCHEMA.COLUMNS`
+        WHERE table_name = '{table_name}'
+        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]:
+        """BigQuery doesn't enforce primary keys; rely on data-driven grain detection."""
+        table_name, schema = parse_bigquery_table_name(table_name, schema)
+
+        pk_query = f"""
+        SELECT column_name
+        FROM `{schema}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE`
+        WHERE table_name = '{table_name}'
+        AND constraint_name LIKE '%PRIMARY%'
+        """
+
+        rows = executor.execute_raw_sql(pk_query).fetchall()
+        return [row[0] for row in rows]
+
+    def render_array_unnest(
+        self,
+        left,
+        right,
+        operator: ComparisonOperator,
+        cte: CTE | UnionCTE | None = None,
+        cte_map: Optional[Dict[str, CTE | UnionCTE]] = None,
+        raise_invalid: bool = False,
+    ):
+        return f"{self.render_expr(left, cte=cte, cte_map=cte_map, raise_invalid=raise_invalid)} {operator.value} unnest({self.render_expr(right, cte=cte, cte_map=cte_map, raise_invalid=raise_invalid)})"
+
+    def generate_partitioned_insert(
+        self,
+        query: ProcessedQueryPersist,
+        recursive: bool,
+        compiled_ctes: list[CompiledCTE],
+    ) -> str:
+        tmp_table = f"tmp__{uuid.uuid4().hex}"
+        final_select = compiled_ctes[-1].statement
+        ctes = compiled_ctes[:-1]
+
+        if not query.partition_by:
+            raise ValueError("partition_by must be set for partitioned inserts.")
+
+        partition_key = query.partition_by
+        target_table = safe_quote(
+            query.output_to.address.location, self.QUOTE_CHARACTER
+        )
+
+        # render intermediate CTEs
+        ctes_sql = ""
+        if ctes:
+            rendered = []
+            for c in ctes:
+                rendered.append(f"{c.name} AS ({c.statement})")
+            ctes_sql = "WITH " + ",\n".join(rendered)
+
+        # create temp table first
+        full_select_with_ctes = (
+            final_select if not ctes_sql else f"{ctes_sql}\n{final_select}"
+        )
+
+        sql_script = PARTITIONED_INSERT_TEMPLATE.render(
+            tmp_table=tmp_table,
+            final_select=full_select_with_ctes,
+            partition_key=partition_key,
+            target_table=target_table,
+            partition_type=self.DATATYPE_MAP[query.partition_types[0]],
+        )
+
+        return sql_script

trilogy/dialect/common.py

@@ -0,0 +1,147 @@
+from typing import Callable
+
+from trilogy.core.constants import UNNEST_NAME
+from trilogy.core.enums import Modifier, UnnestMode
+from trilogy.core.models.build import (
+    BuildComparison,
+    BuildConcept,
+    BuildConditional,
+    BuildFunction,
+    BuildParamaterizedConceptReference,
+    BuildParenthetical,
+)
+from trilogy.core.models.execute import (
+    CTE,
+    InstantiatedUnnestJoin,
+    Join,
+    UnionCTE,
+)
+
+
+def render_unnest(
+    unnest_mode: UnnestMode,
+    quote_character: str,
+    concept: BuildConcept | BuildParamaterizedConceptReference | BuildFunction,
+    render_func: Callable[
+        [BuildConcept | BuildParamaterizedConceptReference | BuildFunction, CTE], str
+    ],
+    cte: CTE,
+):
+    if not isinstance(concept, (BuildConcept, BuildParamaterizedConceptReference)):
+        address = UNNEST_NAME
+    else:
+        address = concept.safe_address
+    if unnest_mode == UnnestMode.CROSS_JOIN:
+        return f"{render_func(concept, cte)} as {quote_character}{address}{quote_character}"
+    elif unnest_mode == UnnestMode.CROSS_JOIN_UNNEST:
+        return f"unnest({render_func(concept, cte)}) as {quote_character}{address}{quote_character}"
+    elif unnest_mode == UnnestMode.PRESTO:
+        return f"unnest({render_func(concept, cte)}) as t({quote_character}{UNNEST_NAME}{quote_character})"
+    elif unnest_mode == UnnestMode.CROSS_JOIN_ALIAS:
+        return f"{render_func(concept, cte)} as unnest_wrapper ({quote_character}{address}{quote_character})"
+    elif unnest_mode == UnnestMode.SNOWFLAKE:
+        # if we don't actually have a join, we're directly unnesting a concept, and we can skip the flatten
+        if not cte.render_from_clause:
+            return f"{render_func(concept, cte)} as unnest_wrapper ( unnest1, unnest2, unnest3, unnest4, {quote_character}{cte.join_derived_concepts[0].safe_address}{quote_character})"
+        # otherwise, flatten the concept for the join
+        return f"flatten({render_func(concept, cte)}) as unnest_wrapper ( unnest1, unnest2, unnest3, unnest4, {quote_character}{cte.join_derived_concepts[0].safe_address}{quote_character})"
+    return f"{render_func(concept, cte)} as {quote_character}{address}{quote_character}"
+
+
+def render_join_concept(
+    name: str,
+    quote_character: str,
+    cte: CTE | UnionCTE,
+    concept: BuildConcept,
+    render_expr,
+    inlined_ctes: set[str],
+    use_map: dict[str, set[str]],
+):
+    if cte.name in inlined_ctes:
+        base = render_expr(concept, cte)
+        return base
+    use_map[name].add(concept.address)
+    return f"{quote_character}{name}{quote_character}.{quote_character}{concept.safe_address}{quote_character}"
+
+
+def render_join(
+    join: Join | InstantiatedUnnestJoin,
+    quote_character: str,
+    render_expr_func: Callable[
+        [
+            BuildConcept
+            | BuildParamaterizedConceptReference
+            | BuildFunction
+            | BuildConditional
+            | BuildComparison
+            | BuildParenthetical,
+            CTE,
+        ],
+        str,
+    ],
+    cte: CTE,
+    use_map: dict[str, set[str]],
+    null_wrapper: Callable[[str, str, list[Modifier]], str],
+    unnest_mode: UnnestMode = UnnestMode.CROSS_APPLY,
+) -> str | None:
+    # {% for key in join.joinkeys %}{{ key.inner }} = {{ key.outer}}{% endfor %}
+    if isinstance(join, InstantiatedUnnestJoin):
+        if unnest_mode == UnnestMode.DIRECT:
+            return None
+        if not cte:
+            raise ValueError("must provide a cte to build an unnest joins")
+        if unnest_mode in (
+            UnnestMode.CROSS_JOIN,
+            UnnestMode.CROSS_JOIN_UNNEST,
+            UnnestMode.CROSS_JOIN_ALIAS,
+            UnnestMode.PRESTO,
+        ):
+            return f"CROSS JOIN {render_unnest(unnest_mode, quote_character, join.object_to_unnest, render_expr_func, cte)}"
+        if unnest_mode == UnnestMode.SNOWFLAKE:
+            return f"LEFT JOIN LATERAL {render_unnest(unnest_mode, quote_character, join.object_to_unnest, render_expr_func, cte)}"
+        return f"FULL JOIN {render_unnest(unnest_mode, quote_character, join.object_to_unnest, render_expr_func, cte)}"
+    # left_name = join.left_name
+    right_name = join.right_name
+    join.quote = quote_character
+    # if cte.quote_address.get(join.right_name, False):
+    #     join.quote = quote_character
+    right_base = join.right_ref
+    base_joinkeys = []
+    if join.joinkey_pairs:
+        base_joinkeys.extend(
+            [
+                null_wrapper(
+                    render_join_concept(
+                        join.get_name(pair.cte),
+                        quote_character,
+                        pair.cte,
+                        pair.left,
+                        render_expr_func,
+                        join.inlined_ctes,
+                        use_map=use_map,
+                    ),
+                    render_join_concept(
+                        right_name,
+                        quote_character,
+                        join.right_cte,
+                        pair.right,
+                        render_expr_func,
+                        join.inlined_ctes,
+                        use_map=use_map,
+                    ),
+                    pair.modifiers
+                    + (pair.left.modifiers or [])
+                    + (pair.right.modifiers or [])
+                    + (join.modifiers or []),
+                )
+                for pair in join.joinkey_pairs
+            ]
+        )
+    if not base_joinkeys:
+        base_joinkeys = ["1=1"]
+
+    joinkeys = " AND ".join(sorted(base_joinkeys))
+    base = f"{join.jointype.value.upper()} JOIN {right_base} on {joinkeys}"
+    if join.condition:
+        base = f"{base} and {render_expr_func(join.condition, cte)}"
+    return base

trilogy/dialect/config.py

@@ -0,0 +1,144 @@
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    try:
+        from pandas import DataFrame
+    except ImportError:
+        DataFrame = Any
+
+
+class DialectConfig:
+    def __init__(self):
+        pass
+
+    def connection_string(self) -> str:
+        raise NotImplementedError
+
+    def create_connect_args(self) -> dict:
+        return {}
+
+    def merge_config(self, other: "DialectConfig") -> "DialectConfig":
+        for key, value in other.__dict__.items():
+            if value is not None:
+                setattr(self, key, value)
+        return self
+
+
+class BigQueryConfig(DialectConfig):
+    def __init__(self, project: str | None = None, client: Any | None = None):
+        self.project = project
+        self.client = client
+
+    def connection_string(self) -> str:
+        return f"bigquery://{self.project}?user_supplied_client=True"
+
+    def create_connect_args(self) -> dict:
+        if not self.client:
+            from google.auth import default
+            from google.cloud import bigquery
+
+            credentials, project = default()
+            self.client = bigquery.Client(credentials=credentials, project=project)
+            self.project = project
+
+        return {"client": self.client}
+
+
+class DuckDBConfig(DialectConfig):
+    def __init__(self, path: str | None = None):
+        self.path = path
+        self.guid = id(self)
+
+    def connection_string(self) -> str:
+        if not self.path:
+            return "duckdb:///:memory:"
+        return f"duckdb:///{self.path}"
+
+
+class PostgresConfig(DialectConfig):
+    def __init__(
+        self, host: str, port: int, username: str, password: str, database: str
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.database = database
+
+    def connection_string(self) -> str:
+        return f"postgresql://{self.username}:{self.password}@{self.host}:{self.port}"
+
+
+class SQLServerConfig(DialectConfig):
+    def __init__(
+        self, host: str, port: int, username: str, password: str, database: str
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.database = database
+
+    def connection_string(self) -> str:
+        return f"sqlserver//{self.username}:{self.password}@{self.host}:{self.port}"
+
+
+class SnowflakeConfig(DialectConfig):
+    def __init__(
+        self,
+        account: str,
+        username: str,
+        password: str,
+        database: str | None = None,
+        schema: str | None = None,
+    ):
+        self.account = account
+        self.username = username
+        self.password = password
+        self.database = database
+        self.schema = schema
+        if self.schema and not self.database:
+            raise ValueError("Setting snowflake schema also requires setting database")
+
+    def connection_string(self) -> str:
+        if self.schema:
+            return f"snowflake://{self.username}:{self.password}@{self.account}/{self.database}/{self.schema}"
+        if self.database:
+            return f"snowflake://{self.username}:{self.password}@{self.account}/{self.database}"
+        return f"snowflake://{self.username}:{self.password}@{self.account}"
+
+
+class PrestoConfig(DialectConfig):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        username: str,
+        password: str,
+        catalog: str,
+        schema: str | None = None,
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.catalog = catalog
+        self.schema = schema
+
+    def connection_string(self) -> str:
+        if self.schema:
+            return f"presto://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"
+        return f"presto://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}"
+
+
+class TrinoConfig(PrestoConfig):
+    def connection_string(self) -> str:
+        if self.schema:
+            return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"
+        return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}"
+
+
+class DataFrameConfig(DuckDBConfig):
+    def __init__(self, dataframes: dict[str, "DataFrame"]):
+        super().__init__()
+        self.dataframes = dataframes

trilogy/dialect/dataframe.py

@@ -0,0 +1,50 @@
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import text
+
+from trilogy.core.models.environment import Environment
+from trilogy.dialect.duckdb import DuckDBDialect
+from trilogy.engine import ExecutionEngine
+
+if TYPE_CHECKING:
+    try:
+        from pandas import DataFrame
+    except ImportError:
+        DataFrame = Any
+
+
+class DataframeDialect(DuckDBDialect):
+    pass
+
+
+class DataframeConnectionWrapper(ExecutionEngine):
+    def __init__(self, engine: ExecutionEngine, dataframes: dict[str, "DataFrame"]):
+        self.engine = engine
+        self.dataframes = dataframes
+        self.connection = None
+
+    def setup(self, env: Environment, connection):
+        self._register_dataframes(env, connection)
+
+    def _register_dataframes(self, env: Environment, connection):
+        for ds in env.datasources.values():
+            if ds.safe_address in self.dataframes:
+                connection.execute(
+                    text("register(:name, :df)"),
+                    {"name": ds.safe_address, "df": self.dataframes[ds.safe_address]},
+                )
+            else:
+                raise ValueError(
+                    f"Dataframe {ds.safe_address} not found in dataframes on connection config, have {self.dataframes.keys()}"
+                )
+        pass
+
+    def add_dataframe(self, name: str, df: "DataFrame", connection, env: Environment):
+        self.dataframes[name] = df
+        self._register_dataframes(env, connection)
+
+    def connect(self) -> Any:
+        return self.engine.connect()
+
+    def dispose(self, close=True):
+        return super().dispose(close)