diff --git a/evalbench/evaluator/db_manager.py b/evalbench/evaluator/db_manager.py index 015a81b8..5bcffda5 100644 --- a/evalbench/evaluator/db_manager.py +++ b/evalbench/evaluator/db_manager.py @@ -1,5 +1,6 @@ from queue import Queue from copy import deepcopy +import os from databases import DB, get_database from util.config import load_db_data_from_csvs, load_setup_scripts from concurrent.futures import ThreadPoolExecutor @@ -9,9 +10,13 @@ def build_db_queue( - core_db: DB, db_name, db_config, setup_config, query_type: str, num_dbs: int + core_db: DB, db_name, db_config, setup_config, query_type: str, + num_dbs: int ): - logging.info(f"Building DB queue (query_type='{query_type}') with {num_dbs} pools for {db_name}...") + logging.info( + f"Building DB queue (query_type='{query_type}') with {num_dbs} " + f"pools for {db_name}..." + ) if query_type == "dql": return _prepare_db_queue_for_dql( core_db, db_name, db_config, setup_config, num_dbs @@ -25,17 +30,24 @@ def build_db_queue( core_db, db_name, db_config, setup_config, num_dbs ) - logging.info(f"Finished building DB queue for query_type '{query_type}' on {db_name}") + logging.info( + f"Finished building DB queue for query_type '{query_type}' on " + f"{db_name}" + ) return Queue[DB]() -def _prepare_db_queue_for_dql(core_db: DB, db_name, db_config, setup_config, num_dbs): +def _prepare_db_queue_for_dql( + core_db: DB, db_name, db_config, setup_config, num_dbs +): + """For DQL, use the same single DB with a user that has only DQL access.""" db_queue = Queue[DB]() dql_db_config = deepcopy(db_config) if setup_config: setup_scripts, data = _get_setup_values( - setup_config, db_name, db_config.get("db_type") + setup_config, db_name, db_config.get("db_type"), + db_config.get("dialect") ) core_db.set_setup_instructions(setup_scripts, data) core_db.resetup_database(False, True) @@ -47,14 +59,20 @@ def _prepare_db_queue_for_dql(core_db: DB, db_name, db_config, setup_config, num return db_queue -def _prepare_db_queue_for_dml(core_db: DB, db_name, db_config, setup_config, num_dbs): - """For DML, use the same single DB with a user that has only DQL / DML access.""" +def _prepare_db_queue_for_dml( + core_db: DB, db_name, db_config, setup_config, num_dbs +): + """For DML, use the same single DB with a user that has only DQL/DML + access. + """ db_queue = Queue[DB]() dml_db_config = deepcopy(db_config) if setup_config: setup_scripts, data = _get_setup_values( - setup_config, db_name, db_config.get("db_type") + setup_config, db_name, db_config.get("db_type"), + db_config.get("dialect") ) + core_db.set_setup_instructions(setup_scripts, data) core_db.resetup_database(False, True) dml_db_config["user_name"] = core_db.get_dml_user() @@ -65,11 +83,14 @@ def _prepare_db_queue_for_dml(core_db: DB, db_name, db_config, setup_config, num return db_queue -def _prepare_db_queue_for_ddl(core_db: DB, db_name, db_config, setup_config, num_dbs): +def _prepare_db_queue_for_ddl( + core_db: DB, db_name, db_config, setup_config, num_dbs +): """For DDL, use the same single DB with a user that has only DDL access.""" if setup_config: setup_scripts, _ = _get_setup_values( - setup_config, db_name, db_config.get("db_type") + setup_config, db_name, db_config.get("db_type"), + db_config.get("dialect") ) core_db.set_setup_instructions(setup_scripts, None) core_db.resetup_database(False, False) @@ -77,13 +98,16 @@ def _prepare_db_queue_for_ddl(core_db: DB, db_name, db_config, setup_config, num if not setup_config: raise ValueError("No Setup Config was provided for DDL") setup_scripts, _ = _get_setup_values( - setup_config, db_name, db_config.get("db_type") + setup_config, db_name, db_config.get("db_type"), + db_config.get("dialect") ) tmp_dbs = core_db.create_tmp_databases(num_dbs) with ThreadPoolExecutor() as executor: create_ddl_tmp_db_p = partial( - _create_ddl_tmp_db, db_config=db_config, setup_scripts=setup_scripts + _create_ddl_tmp_db, db_config=db_config, + setup_scripts=setup_scripts ) + results = executor.map(create_ddl_tmp_db_p, tmp_dbs) for tmp_db in results: db_queue.put(tmp_db) @@ -98,9 +122,36 @@ def _create_ddl_tmp_db(tmp_db, db_config, setup_scripts): return tmp_db -def _get_setup_values(setup_config, db_name: str, db_type: str): +def _get_setup_values( + setup_config, db_name: str, db_type: str, + dialect: Optional[str] = None +): + current_directory = os.getcwd() + setup_dir = os.path.join( + current_directory, setup_config["setup_directory"], db_name + ) + + # 1. Try dialect path if dialect is provided and exists + if dialect: + dialect_path = os.path.join(setup_dir, dialect) + if os.path.isdir(dialect_path): + setup_scripts_dir = os.path.relpath( + dialect_path, current_directory + ) + try: + setup_scripts = load_setup_scripts(setup_scripts_dir) + data = load_db_data_from_csvs( + setup_config["setup_directory"] + "/" + db_name + "/data" + ) + return setup_scripts, data + except Exception: + pass + + # 2. Try db_type path as fallback try: - scripts_path = setup_config["setup_directory"] + "/" + db_name + "/" + db_type + scripts_path = ( + setup_config["setup_directory"] + "/" + db_name + "/" + db_type + ) data_path = setup_config["setup_directory"] + "/" + db_name + "/data" logging.info(f"Loading DB setup files from location: {scripts_path}") @@ -112,5 +163,6 @@ def _get_setup_values(setup_config, db_name: str, db_type: str): return setup_scripts, data except Exception as e: raise FileNotFoundError( - f"Could not find setup files for database {db_name} on {db_type} due to: {e}" + f"Could not find setup files for database {db_name} on " + f"{db_type}/{dialect} due to: {e}" ) diff --git a/evalbench/test/test_db_manager.py b/evalbench/test/test_db_manager.py new file mode 100644 index 00000000..242ba0fe --- /dev/null +++ b/evalbench/test/test_db_manager.py @@ -0,0 +1,89 @@ +import unittest +from unittest.mock import patch, MagicMock +import os +import sys + +# Ensure evalbench is in sys.path +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) + +from evalbench.evaluator.db_manager import _get_setup_values # noqa: E402 + + +class TestDbManagerSetup(unittest.TestCase): + + @patch('evalbench.evaluator.db_manager.load_setup_scripts') + @patch('evalbench.evaluator.db_manager.load_db_data_from_csvs') + @patch('os.path.isdir') + def test_get_setup_values_with_dialect_exists( + self, mock_isdir, mock_load_data, mock_load_scripts + ): + # Setup: dialect directory exists + mock_isdir.side_effect = lambda path: "spanner_gsql" in path + + setup_config = {"setup_directory": "setup"} + db_name = "test_db" + db_type = "spanner" + dialect = "spanner_gsql" + + mock_load_scripts.return_value = (["pre"], ["setup"], ["post"]) + mock_load_data.return_value = {"table1": ["data"]} + + setup_scripts, data = _get_setup_values( + setup_config, db_name, db_type, dialect + ) + + # Assert: loaded from spanner_gsql + mock_load_scripts.assert_called_once_with("setup/test_db/spanner_gsql") + mock_load_data.assert_called_once_with("setup/test_db/data") + self.assertEqual(setup_scripts, (["pre"], ["setup"], ["post"])) + self.assertEqual(data, {"table1": ["data"]}) + + @patch('evalbench.evaluator.db_manager.load_setup_scripts') + @patch('evalbench.evaluator.db_manager.load_db_data_from_csvs') + @patch('os.path.isdir') + def test_get_setup_values_with_dialect_missing_fallback( + self, mock_isdir, mock_load_data, mock_load_scripts + ): + # Setup: dialect directory does not exist, but db_type directory exists + mock_isdir.return_value = False + + setup_config = {"setup_directory": "setup"} + db_name = "test_db" + db_type = "spanner" + dialect = "spanner_gsql" + + mock_load_scripts.return_value = (["pre"], ["setup"], ["post"]) + mock_load_data.return_value = {"table1": ["data"]} + + setup_scripts, data = _get_setup_values( + setup_config, db_name, db_type, dialect + ) + + # Assert: loaded from fallback spanner + mock_load_scripts.assert_called_once_with("setup/test_db/spanner") + mock_load_data.assert_called_once_with("setup/test_db/data") + + @patch('evalbench.evaluator.db_manager.load_setup_scripts') + @patch('evalbench.evaluator.db_manager.load_db_data_from_csvs') + @patch('os.path.isdir') + def test_get_setup_values_no_dialect( + self, mock_isdir, mock_load_data, mock_load_scripts + ): + setup_config = {"setup_directory": "setup"} + db_name = "test_db" + db_type = "spanner" + + mock_load_scripts.return_value = (["pre"], ["setup"], ["post"]) + mock_load_data.return_value = {"table1": ["data"]} + + setup_scripts, data = _get_setup_values( + setup_config, db_name, db_type, None + ) + + # Assert: loaded from spanner + mock_load_scripts.assert_called_once_with("setup/test_db/spanner") + mock_load_data.assert_called_once_with("setup/test_db/data") + + +if __name__ == '__main__': + unittest.main() diff --git a/evalbench/util/setup_databases.py b/evalbench/util/setup_databases.py index fd41173f..009ec460 100755 --- a/evalbench/util/setup_databases.py +++ b/evalbench/util/setup_databases.py @@ -1,12 +1,13 @@ -""" -Utility script to setup databases (schema and data) for an EvalBench experiment natively. +"""Utility script to setup databases (schema and data) natively. -This script parses a standard EvalBench experiment_config YAML file (the same file -you would pass to evalbench.py), extracts the database requirements and setup paths, -and automatically creates and sets up the schemas in the target engines. +This script parses a standard EvalBench experiment_config YAML file +(the same file you would pass to evalbench.py), extracts the database +requirements and setup paths, and automatically creates and sets up +the schemas in the target engines. Example usage: - python3 evalbench/util/setup_databases.py --experiment_config datasets/bird/example_run_config.yaml + python3 evalbench/util/setup_databases.py \ + --experiment_config datasets/bird/example_run_config.yaml """ @@ -14,6 +15,7 @@ from evalbench.databases import get_database from evalbench.dataset.dataset import load_dataset_from_json, flatten_dataset from evalbench.util.config import load_yaml_config +from evalbench.util.flags import EXPERIMENT_CONFIG import sys import os from absl import app @@ -42,20 +44,23 @@ def setup_databases(config_path: str): dialect = db_config.get("dialect", db_type) for db_name_from_dataset in unique_db_names: - db_name = db_name_mappings.get(dialect, "{db_id}").format(db_id=db_name_from_dataset) + db_name = db_name_mappings.get( + dialect, "{db_id}" + ).format(db_id=db_name_from_dataset) print(f"Processing {db_name} for engine {db_type}...") # Get connection wrapper to the specific database core_db = get_database(db_config, db_name) - # Ensure the permanent database exists BEFORE running resetup_database + # Ensure the permanent database exists BEFORE running + # resetup_database core_db.ensure_database_exists(db_name) # Load setup scripts natively from SQL directory setup_config = config try: setup_scripts, data = _get_setup_values( - setup_config, db_name, db_type) + setup_config, db_name, db_type, dialect) except Exception as e: print(f" Failed to load setup values: {e}") continue @@ -70,9 +75,6 @@ def setup_databases(config_path: str): print(f" Failed to setup {db_name} on {db_type}: {e}") -from util.flags import EXPERIMENT_CONFIG - - def main(argv): if len(argv) > 1: raise app.UsageError("Too many command-line arguments.")