Allow per-dialect schema files

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,25 +83,31 @@ 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)
     db_queue = Queue[DB]()
     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}"
         )

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()

evalbench/util/setup_databases.py

@@ -1,19 +1,21 @@
-"""
-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
 """
 
 
 from evalbench.evaluator.db_manager import _get_setup_values
 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.")