ahvn.utils.db.db_utils module

Database configuration utilities for AgentHeaven.

This module provides functions to parse and resolve database configurations similar to how LLM configurations are handled. It supports multiple database providers (SQLite, PostgreSQL, DuckDB, etc.) and generates SQLAlchemy-ready configurations with URLs and hyperparameters.

ahvn.utils.db.db_utils.resolve_db_config(database=None, provider=None, pool=None, **kwargs)

Compile a database configuration dictionary based on the following order of priority: 1. kwargs 2. provider 3. global configuration When a parameter is specified in multiple places, the one with the highest priority is used. When missing, the provider falls back to the default provider.

Parameters:

• database (str, optional) – The database name to use.

• provider (str, optional) – The database provider name to use (e.g., ‘sqlite’, ‘pg’, ‘duckdb’).

• pool (Dict[str, Any], optional) – Pool configuration to override provider defaults.

• **kwargs – Additional parameters to override in the configuration.

Returns:

1. The resolved database configuration dictionary with ‘url’, ‘pool’, and hyperparameters. Connection parameters (dialect, driver, username, password, host, port, database) are used to build the URL and then removed from the final config. 2. The connection parameters dictionary.

Return type:

Tuple[Dict[str, Any], Dict[str, Any]]

ahvn.utils.db.db_utils.create_database_engine(config, conn_args, autocreate=True)

Create a SQLAlchemy engine from the resolved database configuration.

Uses appropriate connection pooling strategy based on dialect: - SQLite: StaticPool (file) or SingletonThreadPool (:memory:) - DuckDB: NullPool (thread-safe, no pooling needed) - PostgreSQL/MySQL/MSSQL: QueuePool with configurable settings

Pool settings are read from conn_args[‘pool’] (set by resolve_db_config from provider config).

Parameters:

• config (Dict[str, Any]) – The database configuration dictionary.

• conn_args (Optional[Dict[str, Any]]) – Connection arguments containing dialect, database, pool config, etc.

• autocreate (bool) – Whether to automatically create the database if it does not exist.

Returns:

A SQLAlchemy engine instance.

Return type:

Engine

Raises:

• ImportError – If SQLAlchemy is not installed.

• ValueError – If required configuration is missing.

ahvn.utils.db.db_utils.create_database(config, engine_kwargs=None)

Create the database if it does not already exist.

This helper supports SQLite (directory creation), PostgreSQL, and MySQL database creation.

Parameters:

• config (Dict[str, Any]) – Database configuration dict containing connection parameters.

• engine_kwargs (Optional[Dict[str, Any]]) – Optional kwargs to pass when creating temporary engines.

Return type:

None

Notes

• The function is best-effort and will log on failure; callers may choose to

  ignore failures by catching exceptions.

ahvn.utils.db.db_utils.split_sqls(queries, dialect='sqlite')

Split a string containing multiple SQL queries into a list.

Parameters:

• queries (str) – The SQL queries to split.

• dialect (str) – The SQL dialect to use (default is “sqlite”).

Returns:

A list of individual SQL queries.

Return type:

List[str]

ahvn.utils.db.db_utils.transpile_sql(query, src_dialect='sqlite', tgt_dialect='sqlite')

Transpile a SQL query from one dialect to another.

Parameters:

• query (str) – The SQL query to transpile.

• src_dialect (str) – The source dialect to transpile from.

• tgt_dialect (str) – The target dialect to transpile to (default is the current dialect).

Returns:

The transpiled query.

Return type:

str

Raises:

• ImportError – If SQLGlot is not installed.

• ValueError – If transpilation fails.

ahvn.utils.db.db_utils.prettify_sql(query, dialect='sqlite', comments=True)

Prettify a SQL query for better readability (identify + strip).

Parameters:

• query (str) – The SQL query to prettify.

• dialect (str) – The SQL dialect to use (default is “sqlite”).

• comments (bool) – Whether to keep comments in the output (default is True).

Returns:

The prettified SQL query. If failed, returns the original query stripped.

Return type:

str

ahvn.utils.db.db_utils.compare_sqls(sql1, sql2, db)

Given two SQL queries, execute them on the provided database and compare their results.

Parameters:

• sql1 (str) – The first SQL query.

• sql2 (str) – The second SQL query.

• db – The database instance with an execute_sql method.

Returns:

True if the results are identical, False otherwise.

Return type:

bool

ahvn.utils.db.db_utils.load_builtin_sql(query_name, dialect='sqlite', **kwargs)

Load SQL query from file and return the query for the current dialect.

Warning

This function uses string formatting (.format(**kwargs)) to inject parameters into the SQL query. This is vulnerable to SQL injection if kwargs contains untrusted user input. Only use this function with trusted input or for internal queries where parameters are controlled. For user-supplied values, prefer using parameterized queries supported by your database driver.

Parameters:

• query_name (str) – Name of the SQL file (without .sql extension).

• dialect (str) – The SQL dialect to use (default is “sqlite”).

• **kwargs – Additional parameters for query formatting.

Returns:

SQL query for the current dialect. None if the query is not found.

Return type:

str

Raises:

FileNotFoundError – If SQL file is not found.

class ahvn.utils.db.db_utils.SQLProcessor(target_dialect)

Bases: object

Handles SQL transpilation and parameter normalization across different database dialects.

This class centralizes all SQL processing logic including: - SQL dialect transpilation via SQLGlot - Parameter format normalization (convert all to :param format) - Cross-database parameter binding support

Parameters:

target_dialect (str)

__init__(target_dialect)

Initialize SQL processor for a target database dialect.

Parameters:

target_dialect (str) – Target database dialect (sqlite, postgres, mysql, duckdb, etc.)

process_query(query, params=None, transpile_from=None)

Process a SQL query and parameters for execution.

Parameters:

• query (str) – Raw SQL query

• params – Query parameters (dict, tuple, list, or None)

• transpile_from (str) – Source dialect to transpile from (if different from target)

Return type:

tuple[str, dict]

Returns:

Tuple of (processed_query, normalized_params)