squirrels 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

squirrels/_command_line.py

@@ -1,5 +1,6 @@
 from argparse import ArgumentParser, _SubParsersAction
-import sys, asyncio, traceback, io
+from pathlib import Path
+import sys, asyncio, traceback, io, subprocess
 
 sys.path.append('.')
 
@@ -7,101 +8,185 @@ from ._version import __version__
 from ._api_server import ApiServer
 from ._initializer import Initializer
 from ._package_loader import PackageLoaderIO
-from .project import SquirrelsProject
-from . import _constants as c
+from ._project import SquirrelsProject
+from . import _constants as c, _utils as u
+from ._compile_prompts import prompt_compile_options
+
+
+def _run_duckdb_cli(project: SquirrelsProject, ui: bool):
+    init_sql = u._read_duckdb_init_sql(datalake_db_path=project._datalake_db_path)
+    
+    target_init_path = None
+    if init_sql:
+        target_init_path = Path(c.TARGET_FOLDER, c.DUCKDB_INIT_FILE)
+        target_init_path.parent.mkdir(parents=True, exist_ok=True)
+        target_init_path.write_text(init_sql)
+    
+    init_args = ["-init", str(target_init_path)] if target_init_path else []
+    command = ['duckdb']
+    if ui:
+        command.append('-ui')
+    if init_args:
+        command.extend(init_args)
+    
+    print("Starting DuckDB CLI with command:")
+    print(f"$ {' '.join(command)}")
+    print()
+    print("To exit the DuckDB CLI, enter '.exit'")
+    print()
+    
+    try:
+        subprocess.run(command, check=True)
+    except FileNotFoundError:
+        print("DuckDB CLI not found. Please install it from: https://duckdb.org/docs/installation/")
+    except subprocess.CalledProcessError:
+        pass # ignore errors that occured on duckdb shell commands
 
 
 def main():
     """
     Main entry point for the squirrels command line utilities.
     """
-    def with_help(parser: ArgumentParser):
-        parser.add_argument('-h', '--help', action="help", help="Show this help message and exit")
-        return parser
-
-    parser = with_help(ArgumentParser(description="Command line utilities from the squirrels python package", add_help=False))
+    # Create a parent parser with common logging options
+    parent_parser = ArgumentParser(add_help=False)
+    parent_parser.add_argument('-h', '--help', action="help", help="Show this help message and exit")
+    parent_parser.add_argument('--log-level', type=str, choices=["DEBUG", "INFO", "WARNING"], help='Level of logging to use. Default is from SQRL_LOGGING__LOG_LEVEL environment variable or INFO.')
+    parent_parser.add_argument('--log-format', type=str, choices=["text", "json"], help='Format of the log records. Default is from SQRL_LOGGING__LOG_FORMAT environment variable or text.')
+    parent_parser.add_argument('--log-to-file', action='store_true', help='Enable logging to file(s) in the "logs/" folder with rotation and retention policies.')
+    
+    parser = ArgumentParser(description="Command line utilities from the squirrels python package", add_help=False, parents=[parent_parser])
     parser.add_argument('-V', '--version', action='store_true', help='Show the version and exit')
-    parser.add_argument('--log-level', type=str, choices=["DEBUG", "INFO", "WARNING"], default="INFO", help='Level of logging to use')
-    parser.add_argument('--log-format', type=str, choices=["text", "json"], default="text", help='Format of the log records')
-    parser.add_argument('--log-file', type=str, default=c.LOGS_FILE, help=f'Name of log file to write to in the "logs/" folder. Default is {c.LOGS_FILE}. If name is empty, then file logging is disabled')
     subparsers = parser.add_subparsers(title='commands', dest='command')
 
     def add_subparser(subparsers: _SubParsersAction, cmd: str, help_text: str):
-        subparser = with_help(subparsers.add_parser(cmd, description=help_text, help=help_text, add_help=False))
+        subparser: ArgumentParser = subparsers.add_parser(cmd, description=help_text, help=help_text, add_help=False, parents=[parent_parser])
         return subparser
 
-    init_parser = add_subparser(subparsers, c.INIT_CMD, 'Initialize a squirrels project')
-    init_parser.add_argument('-o', '--overwrite', action='store_true', help="Overwrite files that already exist")
-    init_parser.add_argument('--core', action='store_true', help='Include all core files')
+    init_parser = add_subparser(subparsers, c.INIT_CMD, 'Create a new squirrels project')
+    
+    init_parser.add_argument('name', nargs='?', type=str, help='The name of the project folder to create. Ignored if --curr-dir is used')
+    init_parser.add_argument('--curr-dir', action='store_true', help='Create the project in the current directory')
+    init_parser.add_argument('--use-defaults', action='store_true', help='Use default values for unspecified options (except project folder name) instead of prompting for input')
     init_parser.add_argument('--connections', type=str, choices=c.CONF_FORMAT_CHOICES, help=f'Configure database connections as yaml (default) or python')
     init_parser.add_argument('--parameters', type=str, choices=c.CONF_FORMAT_CHOICES, help=f'Configure parameters as python (default) or yaml')
-    init_parser.add_argument('--dbview', type=str, choices=c.FILE_TYPE_CHOICES, help='Create database view model as sql (default) or python file')
+    init_parser.add_argument('--build', type=str, choices=c.FILE_TYPE_CHOICES, help='Create build model as sql (default) or python file')
     init_parser.add_argument('--federate', type=str, choices=c.FILE_TYPE_CHOICES, help='Create federated model as sql (default) or python file')
-    init_parser.add_argument('--dashboard', action='store_true', help=f'Include a sample dashboard file')
-    init_parser.add_argument('--auth', action='store_true', help=f'Include the {c.AUTH_FILE} file')
+    init_parser.add_argument('--dashboard', type=str, choices=['y', 'n'], help=f'Include (y) or exclude (n, default) a sample dashboard file')
+    init_parser.add_argument('--admin-password', type=str, help='The password for the admin user. If --use-defaults is used, then a random password is generated')
 
     def with_file_format_options(parser: ArgumentParser):
         help_text = "Create model as sql (default) or python file"
         parser.add_argument('--format', type=str, choices=c.FILE_TYPE_CHOICES, default=c.SQL_FILE_TYPE, help=help_text)
         return parser
     
-    get_file_help_text = "Get a sample file for the squirrels project. If the file name already exists, it will be prefixed with a timestamp."
+    get_file_help_text = "Get a sample file for the squirrels project. If the file name already exists, it will be suffixed with a timestamp."
     get_file_parser = add_subparser(subparsers, c.GET_FILE_CMD, get_file_help_text)
     get_file_subparsers = get_file_parser.add_subparsers(title='file_name', dest='file_name')
-    add_subparser(get_file_subparsers, c.ENV_CONFIG_FILE, f'Get a sample {c.ENV_CONFIG_FILE} file')
+    add_subparser(get_file_subparsers, c.DOTENV_FILE, f'Get sample {c.DOTENV_FILE} and {c.DOTENV_FILE}.example files')
+    add_subparser(get_file_subparsers, c.GITIGNORE_FILE, f'Get a sample {c.GITIGNORE_FILE} file')
     manifest_parser = add_subparser(get_file_subparsers, c.MANIFEST_FILE, f'Get a sample {c.MANIFEST_FILE} file')
     manifest_parser.add_argument("--no-connections", action='store_true', help=f'Exclude the connections section')
     manifest_parser.add_argument("--parameters", action='store_true', help=f'Include the parameters section')
     manifest_parser.add_argument("--dashboards", action='store_true', help=f'Include the dashboards section')
-    add_subparser(get_file_subparsers, c.AUTH_FILE, f'Get a sample {c.AUTH_FILE} file')
+    add_subparser(get_file_subparsers, c.USER_FILE, f'Get a sample {c.USER_FILE} file')
     add_subparser(get_file_subparsers, c.CONNECTIONS_FILE, f'Get a sample {c.CONNECTIONS_FILE} file')
     add_subparser(get_file_subparsers, c.PARAMETERS_FILE, f'Get a sample {c.PARAMETERS_FILE} file')
     add_subparser(get_file_subparsers, c.CONTEXT_FILE, f'Get a sample {c.CONTEXT_FILE} file')
-    with_file_format_options(add_subparser(get_file_subparsers, c.DBVIEW_FILE_STEM, f'Get a sample dbview model file'))
+    add_subparser(get_file_subparsers, c.MACROS_FILE, f'Get a sample {c.MACROS_FILE} file')
+    add_subparser(get_file_subparsers, c.SOURCES_FILE, f'Get a sample {c.SOURCES_FILE} file')
+    with_file_format_options(add_subparser(get_file_subparsers, c.BUILD_FILE_STEM, f'Get a sample build model file'))
+    add_subparser(get_file_subparsers, c.DBVIEW_FILE_STEM, f'Get a sample dbview model file')
     with_file_format_options(add_subparser(get_file_subparsers, c.FEDERATE_FILE_STEM, f'Get a sample federate model file'))
     add_subparser(get_file_subparsers, c.DASHBOARD_FILE_STEM, f'Get a sample dashboard file')
     add_subparser(get_file_subparsers, c.EXPENSES_DB, f'Get the sample SQLite database on expenses')
     add_subparser(get_file_subparsers, c.WEATHER_DB, f'Get the sample SQLite database on weather')
+    add_subparser(get_file_subparsers, c.SEED_CATEGORY_FILE_STEM, f'Get the sample seed files for categories lookup')
+    add_subparser(get_file_subparsers, c.SEED_SUBCATEGORY_FILE_STEM, f'Get the sample seed files for subcategories lookup')
     
-    add_subparser(subparsers, c.DEPS_CMD, f'Load all packages specified in {c.MANIFEST_FILE} (from git)')
+    deps_parser = add_subparser(subparsers, c.DEPS_CMD, f'Load all packages specified in {c.MANIFEST_FILE} (from git)')
 
     compile_parser = add_subparser(subparsers, c.COMPILE_CMD, 'Create rendered SQL files in the folder "./target/compile"')
-    compile_dataset_group = compile_parser.add_mutually_exclusive_group(required=True)
-    compile_dataset_group.add_argument('-d', '--dataset', type=str, help="Select dataset to use for dataset traits. Is required, unless using --all-datasets")
-    compile_dataset_group.add_argument('-D', '--all-datasets', action="store_true", help="Compile models for all datasets. Only required if --dataset is not specified")
+    compile_parser.add_argument('-y', '--yes', action='store_true', help='Disable prompts and assume the default for all settings not overridden by CLI options')
+    compile_parser.add_argument('-c', '--clear', action='store_true', help='Clear the "target/compile/" folder before compiling')
+    compile_scope_group = compile_parser.add_mutually_exclusive_group(required=False)
+    compile_scope_group.add_argument('--buildtime-only', action='store_true', help='Compile only buildtime models')
+    compile_scope_group.add_argument('--runtime-only', action='store_true', help='Compile only runtime models')
     compile_test_set_group = compile_parser.add_mutually_exclusive_group(required=False)
-    compile_test_set_group.add_argument('-t', '--test-set', type=str, help="The selection test set to use. If not specified, default selections are used, unless using --all-test-sets")
-    compile_test_set_group.add_argument('-T', '--all-test-sets', action="store_true", help="Compile models for all selection test sets")
-    compile_parser.add_argument('-s', '--select', type=str, help="Select single model to compile. If not specified, all models for the dataset are compiled. Ignored if using --all-datasets")
-    compile_parser.add_argument('-r', '--runquery', action='store_true', help='Runs all target models, and produce the results as csv files')
+    compile_test_set_group.add_argument('-t', '--test-set', type=str, help="The selection test set to use. If not specified, default selections are used (unless using --all-test-sets). Not applicable for buildtime models")
+    compile_test_set_group.add_argument('-T', '--all-test-sets', action="store_true", help="Compile models for all selection test sets. Not applicable for buildtime models")
+    compile_parser.add_argument('-s', '--select', type=str, help="Select single model to compile. If not specified, all models are compiled")
+    compile_parser.add_argument('-r', '--runquery', action='store_true', help='Run runtime models and write CSV outputs too. Does not apply to buildtime models')
+    
+    build_parser = add_subparser(subparsers, c.BUILD_CMD, 'Build the Virtual Data Lake (VDL) for the project')
+    build_parser.add_argument('-f', '--full-refresh', action='store_true', help='Drop all tables before building')
+    build_parser.add_argument('-s', '--select', type=str, help="Select one static model to build. If not specified, all models are built")
+
+    duckdb_parser = add_subparser(subparsers, c.DUCKDB_CMD, 'Run the duckdb command line tool')
+    duckdb_parser.add_argument('--ui', action='store_true', help='Run the duckdb local UI')
 
     run_parser = add_subparser(subparsers, c.RUN_CMD, 'Run the API server')
+    run_parser.add_argument('--build', action='store_true', help='Build the VDL first (without full refresh) before running the API server')
     run_parser.add_argument('--no-cache', action='store_true', help='Do not cache any api results')
     run_parser.add_argument('--host', type=str, default='127.0.0.1', help="The host to run on")
     run_parser.add_argument('--port', type=int, default=4465, help="The port to run on")
 
     args, _ = parser.parse_known_args()
-    project = SquirrelsProject(log_level=args.log_level, log_format=args.log_format, log_file=args.log_file)
     
     if args.version:
         print(__version__)
     elif args.command == c.INIT_CMD:
-        Initializer(overwrite=args.overwrite).init_project(args)
+        Initializer(project_name=args.name, use_curr_dir=args.curr_dir).init_project(args)
     elif args.command == c.GET_FILE_CMD:
         Initializer().get_file(args)
-    elif args.command == c.DEPS_CMD:
-        PackageLoaderIO.load_packages(project._logger, project._manifest_cfg, reload=True)
-    elif args.command in [c.RUN_CMD, c.COMPILE_CMD]:
+    elif args.command is None:
+        print(f'Command is missing. Enter "squirrels -h" for help.')
+    else:
+        project = SquirrelsProject(load_dotenv_globally=True, log_level=args.log_level, log_format=args.log_format, log_to_file=args.log_to_file)
         try:
-            if args.command == c.RUN_CMD:
+            if args.command == c.DEPS_CMD:
+                PackageLoaderIO.load_packages(project._logger, project._manifest_cfg, reload=True)
+            elif args.command == c.BUILD_CMD:
+                task = project.build(full_refresh=args.full_refresh, select=args.select)
+                asyncio.run(task)
+                print()
+            elif args.command == c.DUCKDB_CMD:
+                _run_duckdb_cli(project, args.ui)
+            elif args.command == c.RUN_CMD:
+                if args.build:
+                    task = project.build(full_refresh=True)
+                    asyncio.run(task)
                 server = ApiServer(args.no_cache, project)
                 server.run(args)
             elif args.command == c.COMPILE_CMD:
-                task = project.compile(
-                    dataset=args.dataset, do_all_datasets=args.all_datasets, selected_model=args.select, test_set=args.test_set, 
-                    do_all_test_sets=args.all_test_sets, runquery=args.runquery
-                )
-                asyncio.run(task)
+                # Derive final options with optional interactive prompts (unless --yes is provided)
+                buildtime_only = args.buildtime_only
+                runtime_only = args.runtime_only
+                test_set = args.test_set
+                do_all_test_sets = args.all_test_sets
+                selected_model = args.select
+
+                try:
+                    if not args.yes:
+                        buildtime_only, runtime_only, test_set, do_all_test_sets, selected_model = prompt_compile_options(
+                            project,
+                            buildtime_only=buildtime_only,
+                            runtime_only=runtime_only,
+                            test_set=test_set,
+                            do_all_test_sets=do_all_test_sets,
+                            selected_model=selected_model,
+                        )
+
+                    task = project.compile(
+                        selected_model=selected_model, test_set=test_set, do_all_test_sets=do_all_test_sets, runquery=args.runquery,
+                        clear=args.clear, buildtime_only=buildtime_only, runtime_only=runtime_only
+                    )
+                    asyncio.run(task)
+                
+                except KeyError:
+                    pass
+            else:
+                print(f'Error: No such command "{args.command}". Enter "squirrels -h" for help.')
+
         except KeyboardInterrupt:
             pass
         except Exception as e:
@@ -112,10 +197,6 @@ def main():
             project._logger.error(err_msg)
         finally:
             project.close()
-    elif args.command is None:
-        print(f'Command is missing. Enter "squirrels -h" for help.')
-    else:
-        print(f'Error: No such command "{args.command}". Enter "squirrels -h" for help.')
 
 
 if __name__ == '__main__':

squirrels/_compile_prompts.py

@@ -0,0 +1,147 @@
+from typing import Tuple, Optional
+import inquirer
+
+from ._project import SquirrelsProject
+from . import _constants as c, _models as m
+
+
+def prompt_compile_options(
+    project: SquirrelsProject,
+    *,
+    buildtime_only: bool,
+    runtime_only: bool,
+    test_set: Optional[str],
+    do_all_test_sets: bool,
+    selected_model: Optional[str],
+) -> Tuple[bool, bool, Optional[str], bool, Optional[str]]:
+    """
+    Interactive prompts to derive compile options when --yes is not provided.
+
+    Returns updated values for:
+    (buildtime_only, runtime_only, test_set, do_all_test_sets, selected_model)
+    """
+    # 1) Scope prompt (skip if explicit flags provided)
+    if not buildtime_only and not runtime_only:
+        scope_answer = inquirer.prompt([
+            inquirer.List(
+                'scope',
+                message='Select scope',
+                choices=[
+                    ('All models', 'all'),
+                    ('Buildtime models only', 'buildtime'),
+                    ('Runtime models only', 'runtime'),
+                ],
+                default='all',
+            )
+        ]) or {}
+        scope_val = scope_answer['scope']
+        if scope_val == 'buildtime':
+            buildtime_only, runtime_only = True, False
+        elif scope_val == 'runtime':
+            buildtime_only, runtime_only = False, True
+        else:
+            buildtime_only, runtime_only = False, False
+
+    # 2) Runtime test set prompt (only if runtime is included and -t/-T not provided)
+    runtime_included, buildtime_included = not buildtime_only, not runtime_only
+    if runtime_included and not test_set and not do_all_test_sets:
+        test_mode_answer = inquirer.prompt([
+            inquirer.List(
+                'ts_mode',
+                message='Runtime selections for parameters, configurables, and user attributes?',
+                choices=[
+                    ('Use default selections', 'default'),
+                    ('Use a custom test set', 'custom'),
+                    ('Use all test sets (including default selections)', 'all'),
+                ],
+                default='default',
+            )
+        ]) or {}
+        ts_mode = test_mode_answer['ts_mode']
+        if ts_mode == 'all':
+            do_all_test_sets = True
+        elif ts_mode == 'custom':
+            # Build available test set list
+            ts_names = list(project._manifest_cfg.selection_test_sets.keys())
+            ts_answer = inquirer.prompt([
+                inquirer.List(
+                    'test_set',
+                    message='Pick a selection test set',
+                    choices=ts_names if ts_names else [c.DEFAULT_TEST_SET_NAME],
+                    default=c.DEFAULT_TEST_SET_NAME,
+                )
+            ]) or {}
+            test_set = ts_answer['test_set']
+
+    # 3) Model selection prompt (only if -s not provided)
+    if not selected_model:
+        # Ask whether to compile all or a specific model
+        model_mode_answer = inquirer.prompt([
+            inquirer.List(
+                'model_mode',
+                message='Compile all models in scope or just a specific model?',
+                choices=[
+                    ('All models', 'all'),
+                    ('Select a specific model', 'one'),
+                ],
+                default='all',
+            )
+        ]) or {}
+        model_mode = model_mode_answer['model_mode']
+
+        if model_mode == 'one':
+            # Build list of runtime query models (dbviews + federates) and build models if included
+            models_dict = project._get_models_dict(always_python_df=False)
+
+            valid_model_types: list[m.ModelType] = []
+            if runtime_included:
+                valid_model_types.extend([m.ModelType.DBVIEW, m.ModelType.FEDERATE])
+            if buildtime_included:
+                valid_model_types.append(m.ModelType.BUILD)
+
+            runtime_names = sorted([
+                (f"({model.model_type.value}) {name}", name) for name, model in models_dict.items()
+                if isinstance(model, m.QueryModel) and model.model_type in valid_model_types
+            ])
+            if runtime_names:
+                model_answer = inquirer.prompt([
+                    inquirer.List(
+                        'selected_model',
+                        message='Pick a runtime model to compile',
+                        choices=runtime_names,
+                    )
+                ]) or {}
+                selected_model = model_answer['selected_model']
+
+    # Tips and equivalent command without prompts
+    def _maybe_quote(val: str) -> str:
+        return f'"{val}"' if any(ch.isspace() for ch in val) else val
+
+    parts = ["sqrl", "compile", "-y"]
+    if buildtime_only:
+        parts.append("--buildtime-only")
+    elif runtime_only:
+        parts.append("--runtime-only")
+    if do_all_test_sets:
+        parts.append("-T")
+    elif test_set:
+        parts.extend(["-t", _maybe_quote(test_set)])
+    if selected_model:
+        parts.extend(["-s", _maybe_quote(selected_model)])
+
+    # Pretty tips block
+    tips_header = " Compile tips "
+    border = "=" * 80
+    print(border)
+    print(tips_header)
+    print("-" * len(border))
+    print("Equivalent command (no prompts):")
+    print(f"  $ {' '.join(parts)}")
+    print()
+    print("You can also:")
+    print("  - Add -c/--clear to clear the 'target/compile/' folder before compiling")
+    print("  - Add -r/--runquery to generate CSVs for runtime models")
+    print(border)
+    print()
+
+    return buildtime_only, runtime_only, test_set, do_all_test_sets, selected_model

squirrels/_connection_set.py

@@ -1,11 +1,11 @@
-from dataclasses import dataclass
-from sqlalchemy import Engine, create_engine
-import time, pandas as pd
+from typing import Any
+from dataclasses import dataclass, field
+from sqlalchemy import Engine
+import time, polars as pl, duckdb
 
 from . import _utils as u, _constants as c, _py_module as pm
-from .arguments.init_time_args import ConnectionsArgs
-from ._environcfg import EnvironConfig
-from ._manifest import ManifestConfig
+from ._arguments.init_time_args import ConnectionsArgs
+from ._manifest import ManifestConfig, ConnectionProperties, ConnectionTypeEnum
 
 
 @dataclass
@@ -16,44 +16,63 @@ class ConnectionSet:
     Attributes:
         _engines: A dictionary of connection name to the corresponding sqlalchemy engine
     """
-    _engines: dict[str, Engine]
+    _connections: dict[str, ConnectionProperties | Any] = field(default_factory=dict)
 
-    def get_engines_as_dict(self):
-        return self._engines.copy()
+    def get_connections_as_dict(self):
+        return self._connections.copy()
     
-    def _get_engine(self, conn_name: str) -> Engine:
+    def get_connection(self, conn_name: str) -> ConnectionProperties | Any:
         try:
-            connection_pool = self._engines[conn_name]
+            connection = self._connections[conn_name]
         except KeyError as e:
             raise u.ConfigurationError(f'Connection name "{conn_name}" was not configured') from e
-        return connection_pool
+        return connection
     
-    def run_sql_query_from_conn_name(self, query: str, conn_name: str, placeholders: dict = {}) -> pd.DataFrame:
-        engine = self._get_engine(conn_name)
+    def run_sql_query_from_conn_name(self, query: str, conn_name: str, placeholders: dict = {}) -> pl.DataFrame:
+        conn = self.get_connection(conn_name)
         try:
-            df = pd.read_sql(query, engine, params=placeholders)
+            if isinstance(conn, ConnectionProperties) and (conn.type == ConnectionTypeEnum.CONNECTORX or conn.type == ConnectionTypeEnum.ADBC):
+                if len(placeholders) > 0:
+                    raise u.ConfigurationError(f"Connection '{conn_name}' is a ConnectorX or ADBC connection, which does not support placeholders")
+                df = pl.read_database_uri(query, conn.uri, engine=conn.type.value)
+            elif isinstance(conn, ConnectionProperties) and conn.type == ConnectionTypeEnum.SQLALCHEMY:
+                with conn.engine.connect() as connection:
+                    df = pl.read_database(query, connection, execute_options={"parameters": placeholders})
+            elif isinstance(conn, ConnectionProperties) and conn.type == ConnectionTypeEnum.DUCKDB:
+                path = conn.uri # .replace("duckdb://", "")
+                dd_conn = duckdb.connect(path)
+                try:
+                    df = dd_conn.sql(query, params=placeholders).pl()
+                finally:
+                    dd_conn.close()
+            else:
+                df = pl.read_database(query, conn, execute_options={"parameters": placeholders}) # type: ignore
             return df
         except Exception as e:
             raise RuntimeError(e) from e
 
     def dispose(self) -> None:
         """
-        Disposes of all the engines in this ConnectionSet
+        Disposes / closes all the connections in this ConnectionSet
         """
-        for pool in self._engines.values():
-            if isinstance(pool, Engine):
-                pool.dispose()
+        for conn in self._connections.values():
+            if isinstance(conn, Engine):
+                conn.dispose()
+            elif isinstance(conn, ConnectionProperties):
+                if conn.type == ConnectionTypeEnum.SQLALCHEMY:
+                    conn.engine.dispose()
+            elif hasattr(conn, 'close'):
+                conn.close()
 
 
 class ConnectionSetIO:
 
     @classmethod
-    def load_conn_py_args(cls, logger: u.Logger, env_cfg: EnvironConfig, manifest_cfg: ManifestConfig) -> ConnectionsArgs:
+    def load_conn_py_args(cls, logger: u.Logger, base_path: str, env_vars: dict[str, str], manifest_cfg: ManifestConfig) -> ConnectionsArgs:
         start = time.time()
         
         proj_vars = manifest_cfg.project_variables.model_dump()
-        env_vars = env_cfg.get_all_env_vars()
-        conn_args = ConnectionsArgs(proj_vars, env_vars, env_cfg.get_credential)
+        conn_args = ConnectionsArgs(base_path, proj_vars, env_vars)
         
         logger.log_activity_time("setting up arguments for connections.py", start)
         return conn_args
@@ -68,13 +87,16 @@ class ConnectionSetIO:
             A ConnectionSet with the DB connections from both squirrels.yml and connections.py
         """
         start = time.time()
-        engines: dict[str, Engine] = {}
+        connections: dict[str, ConnectionProperties | Any] = {}
         
         for config in manifest_cfg.connections.values():
-            engines[config.name] = create_engine(config.url)
+            connections[config.name] = ConnectionProperties(
+                label=config.label, type=config.type, uri=config.uri, sa_create_engine_args=config.sa_create_engine_args
+            )
 
-        pm.run_pyconfig_main(base_path, c.CONNECTIONS_FILE, {"connections": engines, "sqrl": conn_args})
-        conn_set = ConnectionSet(engines)
+        pm.run_pyconfig_main(base_path, c.CONNECTIONS_FILE, {"connections": connections, "sqrl": conn_args})
+
+        conn_set = ConnectionSet(connections)
 
         logger.log_activity_time("creating sqlalchemy engines", start)
         return conn_set