squirrels 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

squirrels/_models.py

@@ -1,21 +1,26 @@
 from __future__ import annotations
-from typing import Union, Optional, Callable, Iterable, Any
+from typing import Optional, Callable, Iterable, Any
 from dataclasses import dataclass, field
+from abc import ABCMeta, abstractmethod
 from enum import Enum
 from pathlib import Path
-import sqlite3, pandas as pd, asyncio, os, shutil
+from sqlalchemy import create_engine, text, Connection
+import asyncio, os, shutil, pandas as pd, json
+import matplotlib.pyplot as plt, networkx as nx
 
 from . import _constants as c, _utils as u, _py_module as pm
 from .arguments.run_time_args import ContextArgs, ModelDepsArgs, ModelArgs
 from ._authenticator import User, Authenticator
 from ._connection_set import ConnectionSetIO
-from ._manifest import ManifestIO, DatasetsConfig
+from ._manifest import ManifestIO, DatasetsConfig, DatasetScope
 from ._parameter_sets import ParameterConfigsSetIO, ParameterSet
+from ._seeds import SeedsIO
 from ._timer import timer, time
 
 class ModelType(Enum):
     DBVIEW = 1
     FEDERATE = 2
+    SEED = 3
 
 class QueryType(Enum):
     SQL = 0
@@ -27,7 +32,7 @@ class Materialization(Enum):
 
 
 @dataclass
-class SqlModelConfig:
+class _SqlModelConfig:
     ## Applicable for dbview models
     connection_name: str
 
@@ -58,63 +63,132 @@ ContextFunc = Callable[[dict[str, Any], ContextArgs], None]
 
 
 @dataclass(frozen=True)
-class RawQuery:
+class _RawQuery(metaclass=ABCMeta):
     pass
 
 @dataclass(frozen=True)
-class RawSqlQuery(RawQuery):
+class _RawSqlQuery(_RawQuery):
     query: str
 
 @dataclass(frozen=True)
-class RawPyQuery(RawQuery):
+class _RawPyQuery(_RawQuery):
     query: Callable[[Any], pd.DataFrame]
     dependencies_func: Callable[[Any], Iterable]
 
 
 @dataclass
-class Query:
+class _Query(metaclass=ABCMeta):
     query: Any
 
 @dataclass
-class WorkInProgress:
+class _WorkInProgress(_Query):
     query: None = field(default=None, init=False)
 
 @dataclass
-class SqlModelQuery(Query):
+class _SqlModelQuery(_Query):
     query: str
-    config: SqlModelConfig
+    config: _SqlModelConfig
 
 @dataclass
-class PyModelQuery(Query):
+class _PyModelQuery(_Query):
     query: Callable[[], pd.DataFrame]
 
 
 @dataclass(frozen=True)
-class QueryFile:
+class _QueryFile:
     filepath: str
     model_type: ModelType
     query_type: QueryType
-    raw_query: RawQuery
+    raw_query: _RawQuery
 
 
 @dataclass
-class Model:
+class _Referable(metaclass=ABCMeta):
     name: str
-    query_file: QueryFile
     is_target: bool = field(default=False, init=False)
-    compiled_query: Optional[Query] = field(default=None, init=False)
 
     needs_sql_table: bool = field(default=False, init=False)
-    needs_pandas: bool = False
+    needs_pandas: bool = field(default=False, init=False)
     result: Optional[pd.DataFrame] = field(default=None, init=False, repr=False)
-    
-    wait_count: int = field(default=0, init=False, repr=False)
-    upstreams: dict[str, Model] = field(default_factory=dict, init=False, repr=False)
-    downstreams: dict[str, Model] = field(default_factory=dict, init=False, repr=False)
 
+    wait_count: int = field(default=0, init=False, repr=False)
     confirmed_no_cycles: bool = field(default=False, init=False)
+    upstreams: dict[str, _Referable] = field(default_factory=dict, init=False, repr=False)
+    downstreams: dict[str, _Referable] = field(default_factory=dict, init=False, repr=False)
+
+    @abstractmethod
+    def get_model_type(self) -> ModelType:
+        pass
+
+    async def compile(
+        self, ctx: dict[str, Any], ctx_args: ContextArgs, placeholders: dict[str, Any], models_dict: dict[str, _Referable], recurse: bool
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def get_terminal_nodes(self, depencency_path: set[str]) -> set[str]:
+        pass
+        
+    def _load_pandas_to_table(self, df: pd.DataFrame, conn: Connection) -> None:
+        df.to_sql(self.name, conn, index=False)
+            
+    def _load_table_to_pandas(self, conn: Connection) -> pd.DataFrame:
+        query = f"SELECT * FROM {self.name}"
+        return pd.read_sql(query, conn)
+    
+    async def _trigger(self, conn: Connection, placeholders: dict = {}) -> None:
+        self.wait_count -= 1
+        if (self.wait_count == 0):
+            await self.run_model(conn, placeholders)
+    
+    @abstractmethod
+    async def run_model(self, conn: Connection, placeholders: dict = {}) -> None:
+        coroutines = []
+        for model in self.downstreams.values():
+            coroutines.append(model._trigger(conn, placeholders))
+        await asyncio.gather(*coroutines)
+    
+    def retrieve_dependent_query_models(self, dependent_model_names: set[str]) -> None:
+        pass
+    
+    def get_max_path_length_to_target(self) -> int:
+        if not hasattr(self, "max_path_len_to_target"):
+            path_lengths = []
+            for child_model in self.downstreams.values():
+                path_lengths.append(child_model.get_max_path_length_to_target()+1)
+            if len(path_lengths) > 0:
+                self.max_path_len_to_target = max(path_lengths)
+            else:
+                self.max_path_len_to_target = 0 if self.is_target else None
+        return self.max_path_len_to_target
 
-    def _add_upstream(self, other: Model) -> None:
+
+@dataclass
+class _Seed(_Referable):
+    result: pd.DataFrame
+
+    def get_model_type(self) -> ModelType:
+        return ModelType.SEED
+
+    def get_terminal_nodes(self, depencency_path: set[str]) -> set[str]:
+        return {self.name}
+    
+    async def run_model(self, conn: Connection, placeholders: dict = {}) -> None:
+        if self.needs_sql_table:
+            await asyncio.to_thread(self._load_pandas_to_table, self.result, conn)
+        await super().run_model(conn, placeholders)
+
+
+@dataclass
+class _Model(_Referable):
+    query_file: _QueryFile
+
+    compiled_query: Optional[_Query] = field(default=None, init=False)
+
+    def get_model_type(self) -> ModelType:
+        return self.query_file.model_type
+
+    def _add_upstream(self, other: _Referable) -> None:
         self.upstreams[other.name] = other
         other.downstreams[self.name] = self
         
@@ -137,17 +211,20 @@ class Model:
             materialized = federate_config.materialized
         return Materialization[materialized.upper()]
     
-    async def _compile_sql_model(self, ctx: dict[str, Any], ctx_args: ContextArgs) -> tuple[SqlModelQuery, set]:
-        assert(isinstance(self.query_file.raw_query, RawSqlQuery))
+    async def _compile_sql_model(
+        self, ctx: dict[str, Any], ctx_args: ContextArgs, placeholders: dict[str, Any]
+    ) -> tuple[_SqlModelQuery, set]:
+        assert(isinstance(self.query_file.raw_query, _RawSqlQuery))
+
         raw_query = self.query_file.raw_query.query
-        
         connection_name = self._get_dbview_conn_name()
         materialized = self._get_materialized()
-        configuration = SqlModelConfig(connection_name, materialized)
+        configuration = _SqlModelConfig(connection_name, materialized)
+        is_placeholder = lambda x: x in placeholders
         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, "config": configuration.set_attribute
+            "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,
+            "config": configuration.set_attribute, "is_param_enabled": ctx_args.param_exists
         }
         dependencies = set()
         if self.query_file.model_type == ModelType.FEDERATE:
@@ -157,16 +234,21 @@ class Model:
             kwargs["ref"] = ref
 
         try:
-            query = await asyncio.to_thread(u.render_string, raw_query, kwargs)
+            query = await asyncio.to_thread(u.render_string, raw_query, **kwargs)
         except Exception as e:
             raise u.FileExecutionError(f'Failed to compile sql model "{self.name}"', e)
         
-        compiled_query = SqlModelQuery(query, configuration)
+        compiled_query = _SqlModelQuery(query, configuration)
         return compiled_query, dependencies
     
-    async def _compile_python_model(self, ctx: dict[str, Any], ctx_args: ContextArgs) -> tuple[PyModelQuery, set]:
-        assert(isinstance(self.query_file.raw_query, RawPyQuery))
-        sqrl_args = ModelDepsArgs(ctx_args.proj_vars, ctx_args.env_vars, ctx_args.user, ctx_args.prms, ctx_args.traits, ctx)
+    async def _compile_python_model(
+        self, ctx: dict[str, Any], ctx_args: ContextArgs, placeholders: dict[str, Any]
+    ) -> tuple[_PyModelQuery, set]:
+        assert(isinstance(self.query_file.raw_query, _RawPyQuery))
+        
+        sqrl_args = ModelDepsArgs(
+            ctx_args.proj_vars, ctx_args.env_vars, ctx_args.user, ctx_args.prms, ctx_args.traits, placeholders, ctx
+        )
         try:
             dependencies = await asyncio.to_thread(self.query_file.raw_query.dependencies_func, sqrl_args)
         except Exception as e:
@@ -175,34 +257,42 @@ class Model:
         dbview_conn_name = self._get_dbview_conn_name()
         connections = ConnectionSetIO.obj.get_engines_as_dict()
         ref = lambda x: self.upstreams[x].result
-        sqrl_args = ModelArgs(ctx_args.proj_vars, ctx_args.env_vars, ctx_args.user, ctx_args.prms, ctx_args.traits, 
-                              ctx, dbview_conn_name, connections, ref, set(dependencies))
+        sqrl_args = ModelArgs(
+            ctx_args.proj_vars, ctx_args.env_vars, ctx_args.user, ctx_args.prms, ctx_args.traits, placeholders, ctx, 
+            dbview_conn_name, connections, dependencies, ref
+        )
             
         def compiled_query():
             try:
-                return self.query_file.raw_query.query(sqrl=sqrl_args)
+                raw_query: _RawPyQuery = self.query_file.raw_query
+                return raw_query.query(sqrl=sqrl_args)
             except Exception as e:
                 raise u.FileExecutionError(f'Failed to run "{c.MAIN_FUNC}" function for python model "{self.name}"', e)
         
-        return PyModelQuery(compiled_query), dependencies
+        return _PyModelQuery(compiled_query), dependencies
 
-    async def compile(self, ctx: dict[str, Any], ctx_args: ContextArgs, models_dict: dict[str, Model], recurse: bool) -> None:
+    async def compile(
+        self, ctx: dict[str, Any], ctx_args: ContextArgs, placeholders: dict[str, Any], models_dict: dict[str, _Referable], recurse: bool
+    ) -> None:
         if self.compiled_query is not None:
             return
         else:
-            self.compiled_query = WorkInProgress()
+            self.compiled_query = _WorkInProgress()
         
         start = time.time()
+
         if self.query_file.query_type == QueryType.SQL:
-            compiled_query, dependencies = await self._compile_sql_model(ctx, ctx_args)
+            compiled_query, dependencies = await self._compile_sql_model(ctx, ctx_args, placeholders)
         elif self.query_file.query_type == QueryType.PYTHON:
-            compiled_query, dependencies = await self._compile_python_model(ctx, ctx_args)
+            compiled_query, dependencies = await self._compile_python_model(ctx, ctx_args, placeholders)
         else:
             raise NotImplementedError(f"Query type not supported: {self.query_file.query_type}")
         
         self.compiled_query = compiled_query
         self.wait_count = len(dependencies)
-        timer.add_activity_time(f"compiling model '{self.name}'", start)
+
+        model_type = self.get_model_type().name.lower()
+        timer.add_activity_time(f"compiling {model_type} model '{self.name}'", start)
         
         if not recurse:
             return 
@@ -211,7 +301,7 @@ class Model:
         coroutines = []
         for dep_model in dep_models:
             self._add_upstream(dep_model)
-            coro = dep_model.compile(ctx, ctx_args, models_dict, recurse)
+            coro = dep_model.compile(ctx, ctx_args, placeholders, models_dict, recurse)
             coroutines.append(coro)
         await asyncio.gather(*coroutines)
     
@@ -234,29 +324,16 @@ class Model:
         
         self.confirmed_no_cycles = True
         return terminal_nodes
-        
-    def _load_pandas_to_table(self, df: pd.DataFrame, conn: sqlite3.Connection) -> None:
-        if u.use_duckdb():
-            conn.execute(f"CREATE TABLE {self.name} AS FROM df")
-        else:
-            df.to_sql(self.name, conn, index=False)
-            
-    def _load_table_to_pandas(self, conn: sqlite3.Connection) -> pd.DataFrame:
-        if u.use_duckdb():
-            return conn.execute(f"FROM {self.name}").df()
-        else:
-            query = f"SELECT * FROM {self.name}"
-            return pd.read_sql(query, conn)
 
-    async def _run_sql_model(self, conn: sqlite3.Connection) -> None:
-        assert(isinstance(self.compiled_query, SqlModelQuery))
+    async def _run_sql_model(self, conn: Connection, placeholders: dict = {}) -> None:
+        assert(isinstance(self.compiled_query, _SqlModelQuery))
         config = self.compiled_query.config
         query = self.compiled_query.query
 
         if self.query_file.model_type == ModelType.DBVIEW:
             def run_sql_query():
                 try:
-                    return ConnectionSetIO.obj.run_sql_query_from_conn_name(query, config.connection_name)
+                    return ConnectionSetIO.obj.run_sql_query_from_conn_name(query, config.connection_name, placeholders)
                 except RuntimeError as e:
                     raise u.FileExecutionError(f'Failed to run dbview sql model "{self.name}"', e)
             
@@ -268,7 +345,7 @@ class Model:
             def create_table():
                 create_query = config.get_sql_for_create(self.name, query)
                 try:
-                    return conn.execute(create_query)
+                    return conn.execute(text(create_query), placeholders)
                 except Exception as e:
                     raise u.FileExecutionError(f'Failed to run federate sql model "{self.name}"', e)
             
@@ -276,8 +353,8 @@ class Model:
             if self.needs_pandas or self.is_target:
                 self.result = await asyncio.to_thread(self._load_table_to_pandas, conn)
     
-    async def _run_python_model(self, conn: sqlite3.Connection) -> None:
-        assert(isinstance(self.compiled_query, PyModelQuery))
+    async def _run_python_model(self, conn: Connection) -> None:
+        assert(isinstance(self.compiled_query, _PyModelQuery))
 
         df = await asyncio.to_thread(self.compiled_query.query)
         if self.needs_sql_table:
@@ -285,92 +362,86 @@ class Model:
         if self.needs_pandas or self.is_target:
             self.result = df
     
-    async def run_model(self, conn: sqlite3.Connection) -> None:
+    async def run_model(self, conn: Connection, placeholders: dict = {}) -> None:
         start = time.time()
+        
         if self.query_file.query_type == QueryType.SQL:
-            await self._run_sql_model(conn)
+            await self._run_sql_model(conn, placeholders)
         elif self.query_file.query_type == QueryType.PYTHON:
             await self._run_python_model(conn)
-        timer.add_activity_time(f"running model '{self.name}'", start)
         
-        coroutines = []
-        for model in self.downstreams.values():
-            coroutines.append(model.trigger(conn))
-        await asyncio.gather(*coroutines)
-    
-    async def trigger(self, conn: sqlite3.Connection) -> None:
-        self.wait_count -= 1
-        if (self.wait_count == 0):
-            await self.run_model(conn)
+        model_type = self.get_model_type().name.lower()
+        timer.add_activity_time(f"running {model_type} model '{self.name}'", start)
+        
+        await super().run_model(conn, placeholders)
     
-    def fill_dependent_model_names(self, dependent_model_names: set[str]) -> None:
+    def retrieve_dependent_query_models(self, dependent_model_names: set[str]) -> None:
         if self.name not in dependent_model_names:
             dependent_model_names.add(self.name)
             for dep_model in self.upstreams.values():
-                dep_model.fill_dependent_model_names(dependent_model_names)
+                dep_model.retrieve_dependent_query_models(dependent_model_names)
 
 
 @dataclass
-class DAG:
+class _DAG:
     dataset: DatasetsConfig
-    target_model: Model
-    models_dict: dict[str, Model]
+    target_model: _Referable
+    models_dict: dict[str, _Referable]
     parameter_set: Optional[ParameterSet] = field(default=None, init=False)
+    placeholders: dict[str, Any] = field(init=False, default_factory=dict)
 
     def apply_selections(
         self, user: Optional[User], selections: dict[str, str], *, updates_only: bool = False, request_version: Optional[int] = None
     ) -> None:
         start = time.time()
         dataset_params = self.dataset.parameters
-        parameter_set = ParameterConfigsSetIO.obj.apply_selections(dataset_params, selections, user, updates_only=updates_only, 
-                                                                   request_version=request_version)
+        parameter_set = ParameterConfigsSetIO.obj.apply_selections(
+            dataset_params, selections, user, updates_only=updates_only, request_version=request_version
+        )
         self.parameter_set = parameter_set
-        timer.add_activity_time(f"applying selections for dataset", start)
+        timer.add_activity_time(f"applying selections for dataset '{self.dataset.name}'", start)
     
     def _compile_context(self, context_func: ContextFunc, user: Optional[User]) -> tuple[dict[str, Any], ContextArgs]:
         start = time.time()
         context = {}
         param_args = ParameterConfigsSetIO.args
         prms = self.parameter_set.get_parameters_as_dict()
-        args = ContextArgs(param_args.proj_vars, param_args.env_vars, user, prms, self.dataset.traits)
+        args = ContextArgs(param_args.proj_vars, param_args.env_vars, user, prms, self.dataset.traits, self.placeholders)
         try:
             context_func(ctx=context, sqrl=args)
         except Exception as e:
-            raise u.FileExecutionError(f'Failed to run {c.CONTEXT_FILE} for dataset "{self.dataset}"', e)
-        timer.add_activity_time(f"running context.py for dataset", start)
+            raise u.FileExecutionError(f'Failed to run {c.CONTEXT_FILE} for dataset "{self.dataset.name}"', e)
+        timer.add_activity_time(f"running context.py for dataset '{self.dataset.name}'", start)
         return context, args
     
     async def _compile_models(self, context: dict[str, Any], ctx_args: ContextArgs, recurse: bool) -> None:
-        await self.target_model.compile(context, ctx_args, self.models_dict, recurse)
+        await self.target_model.compile(context, ctx_args, self.placeholders, self.models_dict, recurse)
     
     def _get_terminal_nodes(self) -> set[str]:
         start = time.time()
         terminal_nodes = self.target_model.get_terminal_nodes(set())
         for model in self.models_dict.values():
             model.confirmed_no_cycles = False
-        timer.add_activity_time(f"validating no cycles in models dependencies", start)
+        timer.add_activity_time(f"validating no cycles in model dependencies", start)
         return terminal_nodes
 
-    async def _run_models(self, terminal_nodes: set[str]) -> None:
-        if u.use_duckdb():
-            import duckdb
-            conn = duckdb.connect()
-        else:
-            conn = sqlite3.connect(":memory:", check_same_thread=False)
+    async def _run_models(self, terminal_nodes: set[str], placeholders: dict = {}) -> None:
+        conn_url = "duckdb:///" if u.use_duckdb() else "sqlite:///?check_same_thread=False"
+        engine = create_engine(conn_url)
         
-        try:
+        with engine.connect() as conn:
             coroutines = []
             for model_name in terminal_nodes:
                 model = self.models_dict[model_name]
-                coroutines.append(model.run_model(conn))
+                coroutines.append(model.run_model(conn, placeholders))
             await asyncio.gather(*coroutines)
-        finally:
-            conn.close()
+        
+        engine.dispose()
     
     async def execute(
         self, context_func: ContextFunc, user: Optional[User], selections: dict[str, str], *, request_version: Optional[int] = None,
         runquery: bool = True, recurse: bool = True
-    ) -> None:
+    ) -> dict[str, Any]:
         recurse = (recurse or runquery)
 
         self.apply_selections(user, selections, request_version=request_version)
@@ -381,17 +452,35 @@ class DAG:
         
         terminal_nodes = self._get_terminal_nodes()
 
+        placeholders = ctx_args._placeholders.copy()
         if runquery:
-            await self._run_models(terminal_nodes)
+            await self._run_models(terminal_nodes, placeholders)
+
+        return placeholders
     
-    def get_all_model_names(self) -> set[str]:
+    def get_all_query_models(self) -> set[str]:
         all_model_names = set()
-        self.target_model.fill_dependent_model_names(all_model_names)
+        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():
+            model_type = model.get_model_type()
+            level = model.get_max_path_length_to_target()
+            if level is not None:
+                G.add_node(model_name, layer=-level, model_type=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
 
 class ModelsIO:
-    raw_queries_by_model: dict[str, QueryFile] 
+    raw_queries_by_model: dict[str, _QueryFile] 
     context_func: ContextFunc
 
     @classmethod
@@ -400,7 +489,6 @@ class ModelsIO:
         cls.raw_queries_by_model = {}
 
         def populate_raw_queries_for_type(folder_path: Path, model_type: ModelType):
-
             def populate_from_file(dp, file):
                 query_type = None
                 filepath = os.path.join(dp, file)
@@ -409,13 +497,13 @@ class ModelsIO:
                     query_type = QueryType.PYTHON
                     module = pm.PyModule(filepath)
                     dependencies_func = module.get_func_or_class(c.DEP_FUNC, default_attr=lambda x: [])
-                    raw_query = RawPyQuery(module.get_func_or_class(c.MAIN_FUNC), dependencies_func)
+                    raw_query = _RawPyQuery(module.get_func_or_class(c.MAIN_FUNC), dependencies_func)
                 elif extension == '.sql':
                     query_type = QueryType.SQL
-                    raw_query = RawSqlQuery(u.read_file(filepath))
+                    raw_query = _RawSqlQuery(u.read_file(filepath))
                 
                 if query_type is not None:
-                    query_file = QueryFile(filepath, model_type, query_type, raw_query)
+                    query_file = _QueryFile(filepath, model_type, query_type, raw_query)
                     if file_stem in cls.raw_queries_by_model:
                         conflicts = [cls.raw_queries_by_model[file_stem].filepath, filepath]
                         raise u.ConfigurationError(f"Multiple models found for '{file_stem}': {conflicts}")
@@ -431,44 +519,98 @@ class ModelsIO:
         federates_path = u.join_paths(c.MODELS_FOLDER, c.FEDERATES_FOLDER)
         populate_raw_queries_for_type(federates_path, ModelType.FEDERATE)
 
-        context_path = u.join_paths(c.PYCONFIG_FOLDER, c.CONTEXT_FILE)
+        context_path = u.join_paths(c.PYCONFIGS_FOLDER, c.CONTEXT_FILE)
         cls.context_func = pm.PyModule(context_path).get_func_or_class(c.MAIN_FUNC, default_attr=lambda x, y: None)
         
-        timer.add_activity_time("loading models and/or context.py", start)
+        timer.add_activity_time("loading files for models and context.py", start)
 
     @classmethod
-    def GenerateDAG(cls, dataset: str, *, target_model_name: Optional[str] = None, always_pandas: bool = False) -> DAG:
-        models_dict = {key: Model(key, val, needs_pandas=always_pandas) for key, val in cls.raw_queries_by_model.items()}
+    def GenerateDAG(cls, dataset: str, *, target_model_name: Optional[str] = None, always_pandas: bool = False) -> _DAG:
+        seeds_dict = SeedsIO.obj.get_dataframes()
+
+        models_dict: dict[str, _Referable] = {key: _Seed(key, df) for key, df in seeds_dict.items()}
+        for key, val in cls.raw_queries_by_model.items():
+            models_dict[key] = _Model(key, val)
+            models_dict[key].needs_pandas = always_pandas
         
         dataset_config = ManifestIO.obj.datasets[dataset]
         target_model_name = dataset_config.model if target_model_name is None else target_model_name
         target_model = models_dict[target_model_name]
         target_model.is_target = True
         
-        return DAG(dataset_config, target_model, models_dict)
+        return _DAG(dataset_config, target_model, models_dict)
     
     @classmethod
-    async def WriteDatasetOutputsGivenTestSet(cls, dataset: str, select: str, test_set: str, runquery: bool, recurse: bool) -> Any:
-        test_set_conf = ManifestIO.obj.selection_test_sets[test_set]
-        user_attributes = test_set_conf.user_attributes
-        selections = test_set_conf.parameters
+    def draw_dag(cls, dag: _DAG, output_folder: Path) -> None:
+        color_map = {ModelType.SEED: "green", ModelType.DBVIEW: "red", ModelType.FEDERATE: "skyblue"}
+
+        G = dag.to_networkx_graph()
+        
+        fig, _ = plt.subplots()
+        pos = nx.multipartite_layout(G, subset_key="layer")
+        colors = [color_map[node[1]] for node in G.nodes(data="model_type")]
+        nx.draw(G, pos=pos, node_shape='^', node_size=1000, node_color=colors, arrowsize=20)
+        
+        y_values = [val[1] for val in pos.values()]
+        scale = max(y_values) - min(y_values) if len(y_values) > 0 else 0
+        label_pos = {key: (val[0], val[1]-0.002-0.1*scale) for key, val in pos.items()}
+        nx.draw_networkx_labels(G, pos=label_pos, font_size=8)
+        
+        fig.tight_layout()
+        plt.margins(x=0.1, y=0.1)
+        plt.savefig(u.join_paths(output_folder, "dag.png"))
+        plt.close(fig)
+
+    @classmethod
+    async def WriteDatasetOutputsGivenTestSet(
+        cls, dataset_conf: DatasetsConfig, select: str, test_set: Optional[str], runquery: bool, recurse: bool
+    ) -> Any:
+        dataset = dataset_conf.name
+        default_test_set, default_test_set_conf = ManifestIO.obj.get_default_test_set(dataset)
+        if test_set is None or test_set == default_test_set:
+            test_set, test_set_conf = default_test_set, default_test_set_conf
+        elif test_set in ManifestIO.obj.selection_test_sets:
+            test_set_conf = ManifestIO.obj.selection_test_sets[test_set]
+        else:
+            raise u.InvalidInputError(f"No test set named '{test_set}' was found when compiling dataset '{dataset}'. The test set must be defined if not default for dataset.")
         
-        username, is_internal = user_attributes.get("username", ""), user_attributes.get("is_internal", False)
-        user_cls: type[User] = Authenticator.get_auth_helper().get_func_or_class("User", default_attr=User)
-        user = user_cls.Create(username, test_set_conf.user_attributes, is_internal=is_internal)
+        error_msg_intro = f"Cannot compile dataset '{dataset}' with test set '{test_set}'."
+        if test_set_conf.datasets is not None and dataset not in test_set_conf.datasets:
+            raise u.InvalidInputError(f"{error_msg_intro}\n Applicable datasets for test set '{test_set}' does not include dataset '{dataset}'.")
         
+        user_attributes = test_set_conf.user_attributes.copy()
+        selections = test_set_conf.parameters.copy()
+        username, is_internal = user_attributes.pop("username", ""), user_attributes.pop("is_internal", False)
+        if test_set_conf.is_authenticated:
+            user_cls: type[User] = Authenticator.get_auth_helper().get_func_or_class("User", default_attr=User)
+            user = user_cls.Create(username, is_internal=is_internal, **user_attributes)
+        elif dataset_conf.scope == DatasetScope.PUBLIC:
+            user = None
+        else:
+            raise u.ConfigurationError(f"{error_msg_intro}\n Non-public datasets require a test set with 'user_attributes' section defined")
+        
+        if dataset_conf.scope == DatasetScope.PRIVATE and not is_internal:
+            raise u.ConfigurationError(f"{error_msg_intro}\n Private datasets require a test set with user_attribute 'is_internal' set to true")
+
+        # always_pandas is set to True for creating CSV files from results (when runquery is True)
         dag = cls.GenerateDAG(dataset, target_model_name=select, always_pandas=True)
-        await dag.execute(cls.context_func, user, selections, runquery=runquery, recurse=recurse)
+        placeholders = await dag.execute(cls.context_func, user, selections, runquery=runquery, recurse=recurse)
         
-        output_folder = u.join_paths(c.TARGET_FOLDER, c.COMPILE_FOLDER, test_set, dataset)
+        output_folder = u.join_paths(c.TARGET_FOLDER, c.COMPILE_FOLDER, dataset, test_set)
         if os.path.exists(output_folder):
             shutil.rmtree(output_folder)
+        os.makedirs(output_folder, exist_ok=True)
+        
+        def write_placeholders() -> None:
+            output_filepath = u.join_paths(output_folder, "placeholders.json")
+            with open(output_filepath, 'w') as f:
+                json.dump(placeholders, f, indent=4)
         
-        def write_model_outputs(model: Model) -> None:
+        def write_model_outputs(model: _Model) -> None:
             subfolder = c.DBVIEWS_FOLDER if model.query_file.model_type == ModelType.DBVIEW else c.FEDERATES_FOLDER
             subpath = u.join_paths(output_folder, subfolder)
             os.makedirs(subpath, exist_ok=True)
-            if isinstance(model.compiled_query, SqlModelQuery):
+            if isinstance(model.compiled_query, _SqlModelQuery):
                 output_filepath = u.join_paths(subpath, model.name+'.sql')
                 query = model.compiled_query.query
                 with open(output_filepath, 'w') as f:
@@ -477,39 +619,50 @@ class ModelsIO:
                 output_filepath = u.join_paths(subpath, model.name+'.csv')
                 model.result.to_csv(output_filepath, index=False)
 
-        all_model_names = dag.get_all_model_names()
+        write_placeholders()
+        all_model_names = dag.get_all_query_models()
         coroutines = [asyncio.to_thread(write_model_outputs, dag.models_dict[name]) for name in all_model_names]
         await asyncio.gather(*coroutines)
-        return dag.target_model.compiled_query.query
+
+        if recurse:
+            cls.draw_dag(dag, output_folder)
+        
+        if isinstance(dag.target_model, _Model):
+            return dag.target_model.compiled_query.query # else return None
 
     @classmethod
     async def WriteOutputs(
-        cls, dataset: Optional[str], select: Optional[str], all_test_sets: bool, test_set: Optional[str], runquery: bool
+        cls, dataset: Optional[str], do_all_datasets: bool, select: Optional[str], test_set: Optional[str], do_all_test_sets: bool, 
+        runquery: bool
     ) -> None:
-        if test_set is None:
-            test_set = ManifestIO.obj.settings.get(c.TEST_SET_DEFAULT_USED_SETTING, c.DEFAULT_TEST_SET_NAME)
-
-        if all_test_sets:
-            test_sets = ManifestIO.obj.selection_test_sets.keys()
-        else:
-            test_sets = [test_set]
+        
+        def get_applicable_test_sets(dataset: str) -> list[str]:
+            applicable_test_sets = []
+            for test_set_name, test_set_config in ManifestIO.obj.selection_test_sets.items():
+                if test_set_config.datasets is None or dataset in test_set_config.datasets:
+                    applicable_test_sets.append(test_set_name)
+            return applicable_test_sets
         
         recurse = True
         dataset_configs = ManifestIO.obj.datasets
-        if dataset is None:
-            selected_models = [(dataset.name, dataset.model) for dataset in dataset_configs.values()]
+        if do_all_datasets:
+            selected_models = [(dataset, dataset.model) for dataset in dataset_configs.values()]
         else:
             if select is None:
                 select = dataset_configs[dataset].model
             else:
                 recurse = False
-            selected_models = [(dataset, select)]
+            selected_models = [(dataset_configs[dataset], select)]
         
         coroutines = []
-        for test_set in test_sets:
-            for dataset, select in selected_models:
-                coroutine = cls.WriteDatasetOutputsGivenTestSet(dataset, select, test_set, runquery, recurse)
-                coroutines.append(coroutine)
+        for dataset_conf, select in selected_models:
+            if do_all_test_sets:
+                for test_set_name in get_applicable_test_sets(dataset_conf.name):
+                    coroutine = cls.WriteDatasetOutputsGivenTestSet(dataset_conf, select, test_set_name, runquery, recurse)
+                    coroutines.append(coroutine)
+            
+            coroutine = cls.WriteDatasetOutputsGivenTestSet(dataset_conf, select, test_set, runquery, recurse)
+            coroutines.append(coroutine)
         
         queries = await asyncio.gather(*coroutines)
         if not recurse and len(queries) == 1 and isinstance(queries[0], str):