Fix evalbench core robustness (#248)

* fix(core): improve robustness and database support

- evalbench.py: Support dict-based dataset_config extraction.
- postgres.py: Enable UNIX domain socket connections for local auth.
- oneshotorchestrator.py: Add database name mapping/overrides for multi-engine evaluations.
- analyzer.py: Fix KeyError and ZeroDivisionError during reporting.

* refactor(analyzer): gracefully handle empty results without short-circuiting

* refactor(analyzer): restore LLM metrics logging

* fix(databases): enhance db registry and driver robustness

- databases/__init__.py: Register SpannerDB and MongoDB in factory.
- mysql.py: Improve connection pooling and Cloud SQL handling.
- sqlite.py: Implement copy-on-write for temporary databases to support file-based datasets.

* fix(databases): enhance db registry and driver robustness

- databases/__init__.py: Register SpannerDB and MongoDB in factory.
- mysql.py: Improve connection pooling and Cloud SQL handling.
- sqlite.py: Implement copy-on-write for temporary databases to support file-based datasets.

* fix(dataset): correct progress calculation for multi-dialect evaluations

* Fix python test failures and pycodestyle warnings

- Updated `evalbench/test/mongodb_test.py` to match the expected data format for `insert_data`.
- Fixed `batch_execute` in `evalbench/databases/spanner.py` to correctly use `self.database.update_ddl` instead of just executing standard queries, as Spanner snapshots do not support DDL commands.
- Updated `evalbench/test/spanner_test.py` to use `batch_execute` for DDL statements (CREATE TABLE, DROP TABLE).
- Modified `.pycodestyle` to ignore `W504` (line break after binary operator), which conflicts with `W503` (line break before binary operator).

* Fix trailing whitespace and formatting errors flagged by pycodestyle

- Used autopep8 and sed to clean up all W291 and W293 errors in the python source.

* Fix comment formatting in spanner.py

---------

Co-authored-by: IsmailMehdi <IsmailMehdi@gmail.com>

Author: aseering

evalbench/client/eval_client.py

@@ -44,7 +44,8 @@ def __init__(self, endpoint: str):
 
             # 4. Composite Credentials
             # Combine the SSL channel with the token-based call credentials
-            composite_creds = grpc.composite_channel_credentials(channel_creds, call_creds)
+            composite_creds = grpc.composite_channel_credentials(
+                channel_creds, call_creds)
             self.channel = grpc.aio.secure_channel(address, composite_creds)
 
         self.stub = eval_service_pb2_grpc.EvalServiceStub(self.channel)

evalbench/databases/__init__.py

@@ -16,7 +16,8 @@ def get_database(db_config, db_name) -> DB:
     #   - It will override the provided default database_name
     #   - This is useful as the default db may be "postgres" or a default only used for setup
     if db_name:
-        db_config["database_name"] = db_name
+        suffix = db_config.get("db_name_suffix", "")
+        db_config["database_name"] = f"{db_name}{suffix}"
 
     if db_config["db_type"] == "postgres":
         return PGDB(db_config)

evalbench/databases/bigquery.py

@@ -68,30 +68,37 @@ def _run_execute(query: str, eval_query: Optional[str] = None, rollback=False):
             error = None
             query_replaced = query.replace("{{dataset}}", self.db_name)
             if eval_query is not None:
-                eval_query_replaced = eval_query.replace("{{dataset}}", self.db_name)
+                eval_query_replaced = eval_query.replace(
+                    "{{dataset}}", self.db_name)
             try:
                 if rollback:
                     try:
                         initial_query = "SELECT 1;"
                         job_config = QueryJobConfig(create_session=True)
-                        init_job = self.client.query(initial_query, job_config=job_config)
+                        init_job = self.client.query(
+                            initial_query, job_config=job_config)
                         init_job.result()
                         session_id = init_job.session_info.session_id
-                        conn_props = [ConnectionProperty(key="session_id", value=session_id)]
+                        conn_props = [ConnectionProperty(
+                            key="session_id", value=session_id)]
 
                         self.client.query(
                             "BEGIN TRANSACTION;",
-                            job_config=QueryJobConfig(connection_properties=conn_props)
+                            job_config=QueryJobConfig(
+                                connection_properties=conn_props)
                         ).result()
 
-                        result = self._execute_queries(query_replaced, job_config=QueryJobConfig(connection_properties=conn_props))
+                        result = self._execute_queries(
+                            query_replaced, job_config=QueryJobConfig(connection_properties=conn_props))
 
                         if eval_query:
-                            eval_result = self._execute_queries(eval_query_replaced, job_config=QueryJobConfig(connection_properties=conn_props))
+                            eval_result = self._execute_queries(
+                                eval_query_replaced, job_config=QueryJobConfig(connection_properties=conn_props))
 
                         self.client.query(
                             "ROLLBACK TRANSACTION;",
-                            job_config=QueryJobConfig(connection_properties=conn_props)
+                            job_config=QueryJobConfig(
+                                connection_properties=conn_props)
                         ).result()
 
                     except Exception as e:
@@ -102,7 +109,8 @@ def _run_execute(query: str, eval_query: Optional[str] = None, rollback=False):
                         if 'session_id' in locals():
                             self.client.query(
                                 "CALL BQ.ABORT_SESSION();",
-                                job_config=QueryJobConfig(connection_properties=conn_props)
+                                job_config=QueryJobConfig(
+                                    connection_properties=conn_props)
                             ).result()
                 if not rollback:
                     result = self._execute_queries(query_replaced)
@@ -113,9 +121,11 @@ def _run_execute(query: str, eval_query: Optional[str] = None, rollback=False):
             except (GoogleAPICallError, Exception) as e:
                 error = str(e)
                 if "resources exceeded" in error:
-                    raise ResourceExhaustedError(f"BigQuery resources exhausted: {e}") from e
+                    raise ResourceExhaustedError(
+                        f"BigQuery resources exhausted: {e}") from e
                 elif "quota exceeded" in error:
-                    raise ResourceExhaustedError(f"BigQuery quota exceeded: {e}") from e
+                    raise ResourceExhaustedError(
+                        f"BigQuery quota exceeded: {e}") from e
                 else:
                     print(error)
 
@@ -140,9 +150,11 @@ def get_metadata(self) -> dict:
         try:
             for table in self.client.list_tables(self.db_name):
                 schema = self.client.get_table(table.reference).schema
-                metadata[table.table_id] = [{"name": f.name, "type": f.field_type} for f in schema]
+                metadata[table.table_id] = [
+                    {"name": f.name, "type": f.field_type} for f in schema]
         except Exception as e:
-            print(f"Error while fetching metadata for dataset '{self.db_name}': {e}")
+            print(
+                f"Error while fetching metadata for dataset '{self.db_name}': {e}")
         return metadata
 
     #####################################################
@@ -155,7 +167,8 @@ def generate_ddl(self, schema: DatabaseSchema) -> List[str]:
         ddl_statements = []
         try:
             for table in schema.tables:
-                columns = ", ".join([f"{col.name} {col.type}" for col in table.columns])
+                columns = ", ".join(
+                    [f"{col.name} {col.type}" for col in table.columns])
                 ddl_statements.append(
                     f"CREATE TABLE `{self.project_id}.{self.db_name}.{table.name}` ({columns})"
                 )
@@ -191,7 +204,8 @@ def drop_all_tables(self):
                     self.client.delete_table(full_table_id)
 
         except Exception as e:
-            raise RuntimeError(f"Failed to drop tables in dataset {self.db_name}: {e}")
+            raise RuntimeError(
+                f"Failed to drop tables in dataset {self.db_name}: {e}")
 
     def _is_float(self, value) -> bool:
         try:
@@ -204,11 +218,13 @@ def _get_column_name_to_type_mapping(self, sql_statements: List[str]) -> Dict[st
         schema_mapping = {}
 
         for statement in sql_statements:
-            table_match = re.search(r'CREATE TABLE\s+`{{dataset}}\.(\w+)`', statement)
+            table_match = re.search(
+                r'CREATE TABLE\s+`{{dataset}}\.(\w+)`', statement)
             if not table_match:
                 continue
             table_name = table_match.group(1)
-            column_section_match = re.search(r'\(\n(.*?)\n\)', statement, re.DOTALL)
+            column_section_match = re.search(
+                r'\(\n(.*?)\n\)', statement, re.DOTALL)
             if not column_section_match:
                 continue
             columns_raw = column_section_match.group(1).split(",\n")

evalbench/databases/bigtable.py

@@ -114,7 +114,8 @@ def get_metadata(self) -> dict:
                     {"name": cf, "type": COLUMN_FAMILY_TYPE} for cf in column_families
                 ]
             except Exception:
-                logging.error(f"Failed to get metadata for table {table.table_id}")
+                logging.error(
+                    f"Failed to get metadata for table {table.table_id}")
         return db_metadata
 
     def generate_ddl(

evalbench/databases/db.py

@@ -91,7 +91,8 @@ def setup_tmp_users(self):
         self.dql_user = "tmp_dql_user_" + generate_key()
         self.dml_user = "tmp_dml_user_" + generate_key()
         self.tmp_user_password = generate_key()
-        self.create_tmp_users(self.dql_user, self.dml_user, self.tmp_user_password)
+        self.create_tmp_users(self.dql_user, self.dml_user,
+                              self.tmp_user_password)
         self.tmp_users.extend([self.dql_user, self.dml_user])
 
     def delete_tmp_users(self, users) -> None:

evalbench/databases/mongodb.py

@@ -149,7 +149,8 @@ def generate_ddl(
         for table in schema.tables:
             ddl.append(f"Collection: {table.name}")
             if table.columns:
-                col_descs = [f"{col.name} ({col.type})" for col in table.columns]
+                col_descs = [
+                    f"{col.name} ({col.type})" for col in table.columns]
                 ddl.append(f"  Fields: {', '.join(col_descs)}")
         return ddl
 

evalbench/databases/mysql.py

@@ -2,6 +2,7 @@
 import sqlparse
 from sqlalchemy import text, MetaData
 from sqlalchemy.engine.base import Connection
+import pymysql
 import logging
 from .db import DB
 from google.cloud.sql.connector import Connector
@@ -47,37 +48,81 @@ class MySQLDB(DB):
 
     def __init__(self, db_config):
         super().__init__(db_config)
-        self.connector = Connector()
+
+        # Auto-deduce use_cloud_sql: format is PROJECT:REGION:INSTANCE (2 colons)
+        self.use_cloud_sql = db_config.get("use_cloud_sql")
+        if self.use_cloud_sql is None:
+            self.use_cloud_sql = (self.db_path.count(":") == 2)
+
+        self.connector = Connector() if self.use_cloud_sql else None
 
         def get_conn():
-            conn = self.connector.connect(
-                self.db_path,
-                "pymysql",
-                user=self.username,
-                password=self.password,
-                db=self.db_name,
-            )
-            return conn
+            """Callable for sqlalchemy 'creator' parameter."""
+            if self.use_cloud_sql:
+                return self.connector.connect(
+                    self.db_path,
+                    "pymysql",
+                    user=self.username,
+                    password=self.password,
+                    db=self.db_name,
+                )
+            else:
+                # Local/Direct connection
+                host = self.db_path
+                port = 3306
+                if ":" in self.db_path:
+                    parts = self.db_path.split(":")
+                    host = parts[0]
+                    port = int(parts[1])
+
+                return pymysql.connect(
+                    host=host,
+                    port=port,
+                    user=self.username,
+                    password=self.password or "",
+                    database=self.db_name
+                )
 
-        def get_engine_args():
-            common_args = {
-                "creator": get_conn,
-                "connect_args": {"command_timeout": 60, "multi_statements": True},
+        def get_engine_config():
+            """Returns (db_url, engine_args)"""
+            args = {
+                "connect_args": {},
             }
+            url = ""
+
+            if self.use_cloud_sql:
+                args["creator"] = get_conn
+                # Cloud SQL needs explicit command_timeout and multi_statements
+                args["connect_args"]["command_timeout"] = 60
+                args["connect_args"]["multi_statements"] = True
+                url = "mysql+pymysql://"
+            else:
+                # Standard local connection via URL
+                # SQLAlchemy parses this URL and loads the driver internally
+                args["connect_args"]["connect_timeout"] = 60
+
+                password_part = f":{self.password}" if self.password else ""
+                url = f"mysql+pymysql://{self.username}{password_part}@{self.db_path}/{self.db_name}"
+
+                password_part = f":{self.password}" if self.password else ""
+                url = f"mysql+pymysql://{self.username}{password_part}@{self.db_path}/{self.db_name}"
+
             if "is_tmp_db" in db_config:
-                common_args["pool_size"] = 1
-                common_args["pool_recycle"] = 300
+                args["pool_size"] = 1
+                args["pool_recycle"] = 300
             else:
-                common_args["pool_size"] = 50
-                common_args["pool_recycle"] = 300
-            return common_args
+                args["pool_size"] = 50
+                args["pool_recycle"] = 300
+            return url, args
 
-        self.engine = sqlalchemy.create_engine("mysql+pymysql://", **get_engine_args())
+        db_url, engine_args = get_engine_config()
+        self.engine = sqlalchemy.create_engine(db_url, **engine_args)
 
     def close_connections(self):
         try:
             self.engine.dispose()
-            self.connector.close()
+            if self.connector:
+                self.connector.close()
         except Exception:
             logging.warning(
                 f"Failed to close connections. This may result in idle unused connections."
@@ -143,7 +188,8 @@ def _run_execute(query: str, eval_query: Optional[str] = None, rollback=False):
                         result = self._execute_queries(connection, query)
 
                         if eval_query:
-                            eval_result = self._execute_queries(connection, eval_query)
+                            eval_result = self._execute_queries(
+                                connection, eval_query)
 
                         if batch_commands and len(batch_commands) > 0:
                             for command in batch_commands:
@@ -181,7 +227,8 @@ def get_metadata(self) -> dict:
                 for table in metadata.tables.values():
                     columns = []
                     for column in table.columns:
-                        columns.append({"name": column.name, "type": str(column.type)})
+                        columns.append(
+                            {"name": column.name, "type": str(column.type)})
                     db_metadata[table.name] = columns
         except Exception:
             pass
@@ -203,7 +250,8 @@ def generate_ddl(
             columns = ", ".join(
                 [f"{column.name} {column.type}" for column in table.columns]
             )
-            create_statements.append(f"CREATE TABLE `{table.name}` ({columns});")
+            create_statements.append(
+                f"CREATE TABLE `{table.name}` ({columns});")
         return create_statements
 
     def create_tmp_database(self, database_name: str):