squirrels 0.5.0b4__py3-none-any.whl → 0.5.1__py3-none-any.whl

squirrels/_models.py

@@ -5,26 +5,26 @@ from abc import ABCMeta, abstractmethod
 from enum import Enum
 from pathlib import Path
 import asyncio, os, re, time, duckdb, sqlglot
-import polars as pl, pandas as pd, networkx as nx
+import polars as pl, pandas as pd
 
 from . import _constants as c, _utils as u, _py_module as pm, _model_queries as mq, _model_configs as mc, _sources as src
 from ._schemas import response_models as rm
 from ._exceptions import FileExecutionError, InvalidInputError
 from ._arguments.run_time_args import ContextArgs, ModelArgs, BuildModelArgs
-from ._auth import BaseUser
+from ._auth import AbstractUser
 from ._connection_set import ConnectionsArgs, ConnectionSet, ConnectionProperties
-from ._manifest import DatasetConfig
+from ._manifest import DatasetConfig, ConnectionTypeEnum
 from ._parameter_sets import ParameterConfigsSet, ParametersArgs, ParameterSet
 
 ContextFunc = Callable[[dict[str, Any], ContextArgs], None]
 
 
 class ModelType(Enum):
+    SEED = "seed"
     SOURCE = "source"
+    BUILD = "build"
     DBVIEW = "dbview"
     FEDERATE = "federate"
-    SEED = "seed"
-    BUILD = "build"
 
 
 @dataclass
@@ -79,15 +79,15 @@ class DataModel(metaclass=ABCMeta):
         self.confirmed_no_cycles = True
         return terminal_nodes
     
-    def _load_duckdb_view_to_python_df(self, conn: duckdb.DuckDBPyConnection, *, use_venv: bool = False) -> pl.LazyFrame:
-        table_name = ("venv." if use_venv else "") + self.name
+    def _load_duckdb_view_to_python_df(self, conn: duckdb.DuckDBPyConnection, *, use_datalake: bool = False) -> pl.LazyFrame:
+        table_name = ("vdl." if use_datalake else "") + self.name
         try:
             return conn.sql(f"FROM {table_name}").pl().lazy()
         except duckdb.CatalogException as e:
             raise u.ConfigurationError(f'Failed to load duckdb table or view "{self.name}" to python dataframe') from e
     
     def _run_sql_query_on_connection(self, connection_name: str, query: str, placeholders: dict = {}) -> pl.DataFrame:
-        self.logger.info(f"Running sql query on connection '{connection_name}': {query}")
+        self.logger.debug(f"Running SQL query on connection '{connection_name}':\n{query}")
         return self.conn_set.run_sql_query_from_conn_name(query, connection_name, placeholders)
     
     async def _trigger(self, conn: duckdb.DuckDBPyConnection, placeholders: dict = {}) -> None:
@@ -133,11 +133,13 @@ class DataModel(metaclass=ABCMeta):
     
     def _create_table_from_df(self, conn: duckdb.DuckDBPyConnection, query_result: pl.LazyFrame | pd.DataFrame):
         local_conn = conn.cursor()
+        # local_conn = conn
         try:
-            local_conn.register("df", query_result)
-            local_conn.execute(f"CREATE OR REPLACE TABLE {self.name} AS SELECT * FROM df")
+            assert query_result is not None
+            local_conn.execute(f"CREATE OR REPLACE TABLE {self.name} AS FROM query_result")
         finally:
             local_conn.close()
+            # pass
         
     def process_pass_through_columns(self, models_dict: dict[str, DataModel]) -> None:
         pass
@@ -172,19 +174,26 @@ class StaticModel(DataModel):
     def _get_result(self, conn: duckdb.DuckDBPyConnection) -> pl.LazyFrame:
         local_conn = conn.cursor()
         try:
-            return self._load_duckdb_view_to_python_df(local_conn, use_venv=True)
+            return self._load_duckdb_view_to_python_df(local_conn, use_datalake=True)
         except Exception as e:
-            raise InvalidInputError(409, f'Dependent data model not found.', f'Model "{self.name}" depends on static data models that cannot be found. Trying building the virtual data environment first.')
+            raise InvalidInputError(409, f'dependent_data_model_not_found', f'Model "{self.name}" depends on static data models that cannot be found. Try building the Virtual Data Lake (VDL) first.')
         finally:
             local_conn.close()
     
     async def run_model(self, conn: duckdb.DuckDBPyConnection, placeholders: dict = {}) -> None:
-        start = time.time()
-
         if (self.needs_python_df or self.is_target) and self.result is None:
+            start = time.time()
+
             self.result = await asyncio.to_thread(self._get_result, conn)
-        
-        self.logger.log_activity_time(f"loading static model '{self.name}'", start)
+            
+            self.logger.log_activity_time(
+                f"loading {self.model_type.value} model '{self.name}' into memory", start, 
+                additional_data={
+                    "activity": "loading static data model into memory",
+                    "model_name": self.name, 
+                    "model_type": self.model_type.value
+                }
+            )
 
         await super().run_model(conn, placeholders)
 
@@ -225,10 +234,18 @@ class Seed(StaticModel):
         start = time.time()
 
         print(f"[{u.get_current_time()}] 🔨 BUILDING: seed model '{self.name}'")
-        await asyncio.to_thread(self._create_table_from_df, conn, self.result)
+        # await asyncio.to_thread(self._create_table_from_df, conn, self.result)
+        self._create_table_from_df(conn, self.result) # without threading
 
         print(f"[{u.get_current_time()}] ✅ FINISHED: seed model '{self.name}'")
-        self.logger.log_activity_time(f"building seed model '{self.name}' to venv", start)
+        self.logger.log_activity_time(
+            f"building seed model '{self.name}' into VDL", start, 
+            additional_data={
+                "activity": "building data model into VDL", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
 
         await super().build_model(conn, full_refresh)
 
@@ -240,24 +257,32 @@ class SourceModel(StaticModel):
     @property
     def model_type(self) -> ModelType:
         return ModelType.SOURCE
+
+    @property
+    def connection_props(self) -> ConnectionProperties:
+        conn_name = self.model_config.get_connection()
+        conn_props = self.conn_set.get_connection(conn_name)
+        if isinstance(conn_props, ConnectionProperties):
+            return conn_props
+        raise u.ConfigurationError(f'Unable to use connection "{conn_name}" for source "{self.name}". Connection "{conn_name}" must be a ConnectionProperties object')
     
     @property
     def is_queryable(self) -> bool:
-        return self.model_config.load_to_duckdb
+        connection_props = self.connection_props
+        return self.model_config.load_to_vdl or connection_props.type == ConnectionTypeEnum.DUCKDB
     
     def _build_source_model(self, conn: duckdb.DuckDBPyConnection, full_refresh: bool) -> None:
         local_conn = conn.cursor()
+        # local_conn = conn
+        
+        local_conn.begin()
         try:
             source = self.model_config
             conn_name = source.get_connection()
             
-            connection_props = self.conn_set.get_connection(conn_name)
-            if isinstance(connection_props, ConnectionProperties):
-                dialect = connection_props.dialect
-                attach_uri = connection_props.attach_uri_for_duckdb
-            else:
-                raise u.ConfigurationError(f'Unable to use connection "{conn_name}" for source "{self.name}". Connection "{conn_name}" must be a ConnectionProperties object')
-            
+            connection_props = self.connection_props
+            dialect = connection_props.dialect
+            attach_uri = connection_props.attach_uri_for_duckdb
             if attach_uri is None:
                 raise u.ConfigurationError(f'Loading to duckdb is not supported for source "{self.name}" since its connection "{conn_name}" uses an unsupported dialect')
             
@@ -269,8 +294,9 @@ class SourceModel(StaticModel):
             new_table_name = self.name
 
             if len(source.columns) == 0:
-                stmt = f"CREATE OR REPLACE TABLE {new_table_name} AS SELECT * FROM db_{conn_name}.{table_name}"
+                stmt = f"CREATE OR REPLACE TABLE {new_table_name} AS FROM db_{conn_name}.{table_name}"
                 u.run_duckdb_stmt(self.logger, local_conn, stmt)
+                local_conn.commit()
                 return
             
             increasing_column = source.update_hints.increasing_column
@@ -297,25 +323,44 @@ class SourceModel(StaticModel):
                 if max_val_of_incr_col is None:
                     recreate_table = True
 
-            insert_cols_clause = source.get_cols_for_insert_stmt()
-            insert_replace_clause = source.get_insert_replace_clause()
-            query = source.get_query_for_insert(dialect, conn_name, table_name, max_val_of_incr_col, full_refresh=recreate_table)
-            stmt = f"INSERT {insert_replace_clause} INTO {new_table_name} ({insert_cols_clause}) {query}"
+            query = source.get_query_for_upsert(dialect, conn_name, table_name, max_val_of_incr_col, full_refresh=recreate_table)
+
+            primary_keys = ", ".join(source.primary_key) if source.primary_key else ""
+            match_condition = f"USING ({primary_keys})" if primary_keys else "ON false"
+            stmt = (
+                f"MERGE INTO {new_table_name} "
+                f"USING ({query}) AS src "
+                f"{match_condition} "
+                f"WHEN MATCHED THEN UPDATE "
+                f"WHEN NOT MATCHED THEN INSERT BY NAME"
+            )
             u.run_duckdb_stmt(self.logger, local_conn, stmt)
+
+            local_conn.commit()
+        
         finally:
             local_conn.close()
+            # pass
 
     async def build_model(self, conn: duckdb.DuckDBPyConnection, full_refresh: bool) -> None:
-        if self.model_config.load_to_duckdb:
+        if self.model_config.load_to_vdl:
             start = time.time()
             print(f"[{u.get_current_time()}] 🔨 BUILDING: source model '{self.name}'")
 
-            await asyncio.to_thread(self._build_source_model, conn, full_refresh)
+            # await asyncio.to_thread(self._build_source_model, conn, full_refresh)
+            self._build_source_model(conn, full_refresh) # without threading
             
             print(f"[{u.get_current_time()}] ✅ FINISHED: source model '{self.name}'")
-            self.logger.log_activity_time(f"building source model '{self.name}' to venv", start)
+            self.logger.log_activity_time(
+                f"building source model '{self.name}' into VDL", start, 
+                additional_data={
+                    "activity": "building data model into VDL", 
+                    "model_name": self.name, 
+                    "model_type": self.model_type.value
+                }
+            )
 
-            await super().build_model(conn, full_refresh)
+        await super().build_model(conn, full_refresh)
         
 
 @dataclass
@@ -338,10 +383,16 @@ class QueryModel(DataModel):
             raise u.ConfigurationError(f'Model "{self.name}" references unknown model "{dependent_model_name}"')
         
         dep_model = models_dict[dependent_model_name]
-        if isinstance(dep_model, SourceModel) and not dep_model.model_config.load_to_duckdb:
-            raise u.ConfigurationError(
-                f'Model "{self.name}" cannot reference source model "{dependent_model_name}" which has load_to_duckdb=False'
-            )
+        if isinstance(dep_model, SourceModel) and not dep_model.model_config.load_to_vdl:
+            # Allow when caller is Build or Federate AND the source connection is duckdb; else error
+            conn_name = dep_model.model_config.get_connection()
+            conn_props = self.conn_set.get_connection(conn_name)
+            is_duckdb_conn = isinstance(conn_props, ConnectionProperties) and conn_props.type == ConnectionTypeEnum.DUCKDB
+            if not is_duckdb_conn:
+                raise u.ConfigurationError(
+                    f'Model "{self.name}" cannot reference source model "{dependent_model_name}". '
+                    'To be referenced by a build or federate model, the source must have load_to_vdl=True or a duckdb connection type.'
+                )
         
         self.model_config.depends_on.add(dependent_model_name)
         return dependent_model_name
@@ -359,7 +410,7 @@ class QueryModel(DataModel):
         is_placeholder = lambda placeholder: placeholder in ctx_args._placeholders_copy
         kwargs = {
             "proj_vars": ctx_args.proj_vars, "env_vars": ctx_args.env_vars, "user": ctx_args.user, "prms": ctx_args.prms, 
-            "traits": ctx_args.traits, "ctx": ctx, "is_placeholder": is_placeholder, "set_placeholder": ctx_args.set_placeholder,
+            "configurables": ctx_args.configurables, "ctx": ctx, "is_placeholder": is_placeholder, "set_placeholder": ctx_args.set_placeholder,
             "param_exists": ctx_args.param_exists
         }
         return kwargs
@@ -433,7 +484,7 @@ class QueryModel(DataModel):
     def _log_sql_to_run(self, sql: str, placeholders: dict[str, Any]) -> None:
         log_msg = f"SQL to run for model '{self.name}':\n{sql}"
         log_msg += f"\n\n(with placeholders: {placeholders})"
-        self.logger.info(log_msg)
+        self.logger.debug(log_msg)
 
 
 @dataclass
@@ -458,11 +509,11 @@ class DbviewModel(QueryModel):
             if source_model.model_config.get_connection() != self.model_config.get_connection():
                 raise u.ConfigurationError(f'Dbview "{self.name}" references source "{source_name}" with different connection')
             
-            # Check if the source model has load_to_duckdb=False but this dbview has translate_to_duckdb=True
-            if not source_model.model_config.load_to_duckdb and self.model_config.translate_to_duckdb:
+            # Check if the source model has load_to_vdl=False but this dbview has translate_to_duckdb=True
+            if not source_model.model_config.load_to_vdl and self.model_config.translate_to_duckdb:
                 raise u.ConfigurationError(
                     f'Dbview "{self.name}" with translate_to_duckdb=True cannot reference source "{source_name}" '
-                    f'which has load_to_duckdb=False'
+                    f'which has load_to_vdl=False'
                 )
                 
             self.model_config.depends_on.add(source_name)
@@ -475,10 +526,11 @@ class DbviewModel(QueryModel):
 
     def _get_duckdb_query(self, read_dialect: str, query: str) -> str:
         kwargs = {
-            "source": lambda source_name: "venv." + source_name
+            "source": lambda source_name: "vdl." + source_name
         }
         compiled_query = self._get_compiled_sql_query_str(query, kwargs)
-        return sqlglot.transpile(compiled_query, read=read_dialect, write="duckdb")[0]
+        duckdb_query = sqlglot.transpile(compiled_query, read=read_dialect, write="duckdb", pretty=True)[0]
+        return "-- translated to duckdb\n" + duckdb_query
     
     def _compile_sql_model(self, kwargs: dict[str, Any]) -> mq.SqlModelQuery:
         compiled_query_str = self._get_compiled_sql_query_str(self.query_file.raw_query, kwargs)
@@ -487,15 +539,20 @@ class DbviewModel(QueryModel):
         connection_props = self.conn_set.get_connection(connection_name)
         
         if self.model_config.translate_to_duckdb and isinstance(connection_props, ConnectionProperties):
-            macros = { 
-                "source": lambda source_name: "venv." + source_name 
+            # Forbid translate_to_duckdb when dbview connection is duckdb
+            if connection_props.type == ConnectionTypeEnum.DUCKDB:
+                raise u.ConfigurationError(
+                    f'Dbview "{self.name}" has translate_to_duckdb=True but its connection is duckdb. Use a federate model instead.'
+                )
+            macros = {
+                "source": lambda source_name: "vdl." + source_name
             }
             compiled_query2 = self._get_compiled_sql_query_str(compiled_query_str, macros)
             compiled_query_str = self._get_duckdb_query(connection_props.dialect, compiled_query2)
             is_duckdb = True
         else:
-            macros = { 
-                "source": lambda source_name: self.sources[source_name].get_table() 
+            macros = {
+                "source": lambda source_name: self.sources[source_name].get_table()
             }
             compiled_query_str = self._get_compiled_sql_query_str(compiled_query_str, macros)
             is_duckdb = False
@@ -516,7 +573,14 @@ class DbviewModel(QueryModel):
         kwargs = self._get_compile_sql_model_args(ctx, ctx_args, models_dict)
         self.compiled_query = self._compile_sql_model(kwargs)
         
-        self.logger.log_activity_time(f"compiling dbview model '{self.name}'", start)
+        self.logger.log_activity_time(
+            f"compiling dbview model '{self.name}'", start, 
+            additional_data={
+                "activity": "compiling data model", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
     
     async def _run_sql_model(self, conn: duckdb.DuckDBPyConnection, placeholders: dict = {}) -> None:
         assert self.compiled_query is not None
@@ -532,7 +596,7 @@ class DbviewModel(QueryModel):
                         self.logger.info(f"Running dbview '{self.name}' on duckdb")
                         return local_conn.sql(query, params=placeholders).pl()
                     except duckdb.CatalogException as e:
-                        raise InvalidInputError(409, f'Dependent data model not found.', f'Model "{self.name}" depends on static data models that cannot be found. Trying building the virtual data environment first.')
+                        raise InvalidInputError(409, f'dependent_data_model_not_found', f'Model "{self.name}" depends on static data models that cannot be found. Try building the Virtual Data Lake (VDL) first.')
                     except Exception as e:
                         raise RuntimeError(e)
                     finally:
@@ -552,7 +616,14 @@ class DbviewModel(QueryModel):
         
         await self._run_sql_model(conn, placeholders)
         
-        self.logger.log_activity_time(f"running dbview model '{self.name}'", start)
+        self.logger.log_activity_time(
+            f"running dbview model '{self.name}'", start, 
+            additional_data={
+                "activity": "running data model", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
         
         await super().run_model(conn, placeholders)
     
@@ -574,8 +645,16 @@ class FederateModel(QueryModel):
         
         def ref(dependent_model_name: str) -> str:
             dependent_model = self._ref_for_sql(dependent_model_name, models_dict)
-            prefix = "venv." if isinstance(models_dict[dependent_model], (SourceModel, BuildModel)) else ""
-            return prefix + dependent_model
+            dep = models_dict[dependent_model]
+            if isinstance(dep, BuildModel):
+                return "vdl." + dependent_model
+            if isinstance(dep, SourceModel):
+                if dep.model_config.load_to_vdl:
+                    return "vdl." + dependent_model
+                conn_name = dep.model_config.get_connection()
+                table_name = dep.model_config.get_table()
+                return f"db_{conn_name}.{table_name}"
+            return dependent_model
         
         kwargs["ref"] = ref
         return kwargs
@@ -629,7 +708,14 @@ class FederateModel(QueryModel):
         else:
             raise NotImplementedError(f"Query type not supported: {self.query_file.__class__.__name__}")
         
-        self.logger.log_activity_time(f"compiling federate model '{self.name}'", start)
+        self.logger.log_activity_time(
+            f"compiling federate model '{self.name}'", start, 
+            additional_data={
+                "activity": "compiling data model", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
         
         if not recurse:
             return 
@@ -658,10 +744,13 @@ class FederateModel(QueryModel):
                 try:
                     return local_conn.execute(create_query, existing_placeholders)
                 except duckdb.CatalogException as e:
-                    raise InvalidInputError(409, f'Dependent data model not found.', f'Model "{self.name}" depends on static data models that cannot be found. Trying building the virtual data environment first.')
+                    if self.name == "__fake_target":
+                        raise InvalidInputError(409, "invalid_sql_query", f"Provided SQL query depends on static data models that cannot be found. Try building the Virtual Data Lake (VDL) first.")
+                    else:
+                        raise InvalidInputError(409, f'dependent_data_model_not_found', f'Model "{self.name}" depends on static data models that cannot be found. Try building the Virtual Data Lake (VDL) first.')
                 except Exception as e:
                     if self.name == "__fake_target":
-                        raise InvalidInputError(400, "Invalid SQL query", f"Failed to run provided SQL query")
+                        raise InvalidInputError(400, "invalid_sql_query", f"Failed to run provided SQL query")
                     else:
                         raise FileExecutionError(f'Failed to run federate sql model "{self.name}"', e) from e
             
@@ -688,7 +777,14 @@ class FederateModel(QueryModel):
         else:
             raise NotImplementedError(f"Query type not supported: {self.query_file.__class__.__name__}")
         
-        self.logger.log_activity_time(f"running federate model '{self.name}'", start)
+        self.logger.log_activity_time(
+            f"running federate model '{self.name}'", start, 
+            additional_data={
+                "activity": "running data model", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
         
         await super().run_model(conn, placeholders)
     
@@ -718,7 +814,12 @@ class BuildModel(StaticModel, QueryModel):
         }
         
         def ref_for_build(dependent_model_name: str) -> str:
-            dependent_model = self._ref_for_sql(dependent_model_name, dict(models_dict))
+            dependent_model = self._ref_for_sql(dependent_model_name, models_dict)
+            dep = models_dict[dependent_model]
+            if isinstance(dep, SourceModel) and not dep.model_config.load_to_vdl:
+                conn_name = dep.model_config.get_connection()
+                table_name = dep.model_config.get_table()
+                return f"db_{conn_name}.{table_name}"
             return dependent_model
         
         kwargs["ref"] = ref_for_build
@@ -771,7 +872,14 @@ class BuildModel(StaticModel, QueryModel):
         else:
             raise NotImplementedError(f"Query type not supported: {self.query_file.__class__.__name__}")
         
-        self.logger.log_activity_time(f"compiling build model '{self.name}'", start)
+        self.logger.log_activity_time(
+            f"compiling build model '{self.name}'", start, 
+            additional_data={
+                "activity": "compiling data model", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
         
         dependencies = self.model_config.depends_on
         self.wait_count_for_build = len(dependencies)
@@ -786,14 +894,17 @@ class BuildModel(StaticModel, QueryModel):
         def create_table():
             create_query = self.model_config.get_sql_for_build(self.name, query)
             local_conn = conn.cursor()
+            # local_conn = conn
             try:
                 return u.run_duckdb_stmt(self.logger, local_conn, create_query, model_name=self.name)
             except Exception as e:
                 raise FileExecutionError(f'Failed to build static sql model "{self.name}"', e) from e
             finally:
                 local_conn.close()
+                # pass
         
-        await asyncio.to_thread(create_table)
+        # await asyncio.to_thread(create_table)
+        create_table() # without threading
 
     async def _build_python_model(self, compiled_query: mq.PyModelQuery, conn: duckdb.DuckDBPyConnection) -> None:
         query_result = await asyncio.to_thread(compiled_query.query)
@@ -801,7 +912,8 @@ class BuildModel(StaticModel, QueryModel):
             query_result = pl.from_pandas(query_result).lazy()
         if self.needs_python_df_for_build:
             self.result = query_result.lazy()
-        await asyncio.to_thread(self._create_table_from_df, conn, query_result)
+        # await asyncio.to_thread(self._create_table_from_df, conn, query_result)
+        self._create_table_from_df(conn, query_result) # without threading
 
     async def build_model(self, conn: duckdb.DuckDBPyConnection, full_refresh: bool) -> None:
         start = time.time()
@@ -814,24 +926,33 @@ class BuildModel(StaticModel, QueryModel):
             def load_df(conn: duckdb.DuckDBPyConnection, dep_model: DataModel):
                 if dep_model.result is None:
                     local_conn = conn.cursor()
+                    # local_conn = conn
                     try:
                         dep_model.result = dep_model._load_duckdb_view_to_python_df(local_conn)
                     finally:
                         local_conn.close()
+                        # pass
                 
             coroutines = []
             for dep_model in self.upstreams_for_build.values():
                 coro = asyncio.to_thread(load_df, conn, dep_model)
                 coroutines.append(coro)
             await u.asyncio_gather(coroutines)
-            
+
             # Then run the model's Python function to build the model
             await self._build_python_model(self.compiled_query, conn)
         else:
             raise NotImplementedError(f"Query type not supported: {self.query_file.__class__.__name__}")
         
         print(f"[{u.get_current_time()}] ✅ FINISHED: build model '{self.name}'")
-        self.logger.log_activity_time(f"building static build model '{self.name}'", start)
+        self.logger.log_activity_time(
+            f"building static build model '{self.name}' into VDL", start, 
+            additional_data={
+                "activity": "building data model into VDL", 
+                "model_name": self.name, 
+                "model_type": self.model_type.value
+            }
+        )
         
         await super().build_model(conn, full_refresh)
 
@@ -841,7 +962,7 @@ class DAG:
     dataset: DatasetConfig | None
     target_model: DataModel
     models_dict: dict[str, DataModel]
-    duckdb_filepath: str = field(default="")
+    datalake_db_path: str | None = field(default=None)
     logger: u.Logger = field(default_factory=lambda: u.Logger(""))
     parameter_set: ParameterSet | None = field(default=None, init=False) # set in apply_selections
     placeholders: dict[str, Any] = field(init=False, default_factory=dict)
@@ -850,36 +971,50 @@ class DAG:
         return f" for dataset '{self.dataset.name}'" if self.dataset else ""
     
     def compile_build_models(self, conn_args: ConnectionsArgs) -> None:
-        static_models: dict[str, StaticModel] = {k: v for k, v in self.models_dict.items() if isinstance(v, StaticModel)}
+        static_models: dict[str, StaticModel] = {
+            k: v for k, v in self.models_dict.items() if isinstance(v, StaticModel)
+        }
         for model in static_models.values():
             if isinstance(model, BuildModel):
                 model.compile_for_build(conn_args, static_models)
 
     def apply_selections(
-        self, param_cfg_set: ParameterConfigsSet, user: BaseUser | None, selections: dict[str, str]
+        self, param_cfg_set: ParameterConfigsSet, user: AbstractUser, selections: dict[str, str]
     ) -> None:
         start = time.time()
+        
         dataset_params = self.dataset.parameters if self.dataset else None
         parameter_set = param_cfg_set.apply_selections(dataset_params, selections, user)
         self.parameter_set = parameter_set
         msg_extension = self._get_msg_extension()
-        self.logger.log_activity_time("applying selections" + msg_extension, start)
+
+        dataset_name = self.dataset.name if self.dataset else None
+        self.logger.log_activity_time(
+            "applying selections" + msg_extension, start, 
+            additional_data={"activity": "applying selections", "dataset_name": dataset_name}
+        )
     
     def _compile_context(
-        self, param_args: ParametersArgs, context_func: ContextFunc, user: BaseUser | None, default_traits: dict[str, Any]
+        self, param_args: ParametersArgs, context_func: ContextFunc, user: AbstractUser, configurables: dict[str, str]
     ) -> tuple[dict[str, Any], ContextArgs]:
         start = time.time()
+
         context = {}
         assert isinstance(self.parameter_set, ParameterSet)
         prms = self.parameter_set.get_parameters_as_dict()
-        traits = self.dataset.traits if self.dataset else default_traits
-        args = ContextArgs(param_args, user, prms, traits)
+        args = ContextArgs(param_args, user, prms, configurables)
         msg_extension = self._get_msg_extension()
+        
         try:
             context_func(context, args)
         except Exception as e:
             raise FileExecutionError(f'Failed to run {c.CONTEXT_FILE}' + msg_extension, e) from e
-        self.logger.log_activity_time("running context.py" + msg_extension, start)
+        
+        dataset_name = self.dataset.name if self.dataset else None
+        self.logger.log_activity_time(
+            "running context.py" + msg_extension, start, 
+            additional_data={"activity": "running context.py", "dataset_name": dataset_name}
+        )
         return context, args
     
     def _compile_models(self, context: dict[str, Any], ctx_args: ContextArgs, recurse: bool) -> None:
@@ -890,29 +1025,25 @@ class DAG:
         terminal_nodes = self.target_model.get_terminal_nodes(set())
         for model in self.models_dict.values():
             model.confirmed_no_cycles = False
-        self.logger.log_activity_time(f"validating no cycles in model dependencies", start)
+        self.logger.log_activity_time("validating no cycles in model dependencies", start)
         return terminal_nodes
 
+    def _attach_connections_with_type_duckdb(self, conn: duckdb.DuckDBPyConnection) -> None:
+        for conn_name, connection in self.target_model.conn_set.get_connections_as_dict().items():
+            if not isinstance(connection, ConnectionProperties):
+                continue
+            attach_uri = connection.attach_uri_for_duckdb
+            if attach_uri is None:
+                continue
+            attach_stmt = f"ATTACH IF NOT EXISTS '{attach_uri}' AS db_{conn_name} (READ_ONLY)"
+            u.run_duckdb_stmt(self.logger, conn, attach_stmt, redacted_values=[attach_uri])
+
     async def _run_models(self) -> None:
         terminal_nodes = self._get_terminal_nodes()
 
-        # create an empty duckdb venv file if it does not exist
+        conn = u.create_duckdb_connection(datalake_db_path=self.datalake_db_path)
         try:
-            conn = duckdb.connect(self.duckdb_filepath)
-            conn.close()
-        except duckdb.IOException as e:
-            # unable to create duckdb venv file means it's in use and already exists
-            # do not throw error here since attaching in read-only mode later may still work
-            pass
-        
-        conn = u.create_duckdb_connection()
-        try:
-            read_only = "(READ_ONLY)" if self.duckdb_filepath else ""
-            try:
-                conn.execute(f"ATTACH '{self.duckdb_filepath}' AS venv {read_only}")
-            except duckdb.IOException as e:
-                self.logger.warning(f"Unable to attach to duckdb venv file: {self.duckdb_filepath}")
-                raise e
+            self._attach_connections_with_type_duckdb(conn)
             
             coroutines = []
             for model_name in terminal_nodes:
@@ -924,14 +1055,14 @@ class DAG:
             conn.close()
     
     async def execute(
-        self, param_args: ParametersArgs, param_cfg_set: ParameterConfigsSet, context_func: ContextFunc, user: BaseUser | None, selections: dict[str, str], 
-        *, runquery: bool = True, recurse: bool = True, default_traits: dict[str, Any] = {}
+        self, param_args: ParametersArgs, param_cfg_set: ParameterConfigsSet, context_func: ContextFunc, user: AbstractUser, selections: dict[str, str], 
+        *, runquery: bool = True, recurse: bool = True, configurables: dict[str, str] = {}
     ) -> None:
         recurse = (recurse or runquery)
 
         self.apply_selections(param_cfg_set, user, selections)
 
-        context, ctx_args = self._compile_context(param_args, context_func, user, default_traits)
+        context, ctx_args = self._compile_context(param_args, context_func, user, configurables)
 
         self._compile_models(context, ctx_args, recurse)
         
@@ -946,21 +1077,6 @@ class DAG:
         self.target_model.retrieve_dependent_query_models(all_model_names)
         return all_model_names
     
-    def to_networkx_graph(self) -> nx.DiGraph:
-        G = nx.DiGraph()
-
-        for model_name, model in self.models_dict.items():
-            level = model.get_max_path_length_to_target()
-            if level is not None:
-                G.add_node(model_name, layer=-level, model_type=model.model_type)
-        
-        for model_name in G.nodes:
-            model = self.models_dict[model_name]
-            for dep_model_name in model.downstreams:
-                G.add_edge(model_name, dep_model_name)
-        
-        return G
-    
     def get_all_data_models(self) -> list[rm.DataModelItem]:
         data_models = []
         for model_name, model in self.models_dict.items():