Adding alloydb omni as a separate db_type

Author: MohammadrezaPourreza

datasets/bat/db_configs/alloydb.yaml

@@ -3,9 +3,6 @@ dialect: postgres
 database_name: <your-db-name>
 database_path: <your-db-path> # should be in the following format: 'projects/{PROJECT_ID}/locations/{REGION}/clusters/{CLUSTER_NAME}/instances/{INSTANCE_NAME}'
 max_executions_per_minute: 180
-host: <your-host> # usef for connecting to omni
-port: <your-port> # usef for connecting to omni, e.g. 5432
 user_name: <your-username>
 password: <your-password>
-nl_config: <nl-config-name> # name of the nl config that will be used for query generation. This should be equal to the name used in the database setup.
-omni: false # set to true if you want to use omni for the database connection
+nl_config: <nl-config-name> # name of the nl config that will be used for query generation. This should be equal to the name used in the database setup.

datasets/bat/run_config.yaml

@@ -5,6 +5,8 @@
 from .db import DB
 from .bigquery import BQDB
 from .alloydb import AlloyDB
+from .alloydb_omni import AlloyDBOmni
+
 
 
 def get_database(db_config, db_name) -> DB:
@@ -26,4 +28,6 @@ def get_database(db_config, db_name) -> DB:
         return BQDB(db_config)
     if db_config["db_type"] == "alloydb":
         return AlloyDB(db_config)
+    if db_config["db_type"] == "alloydb_omni":
+        return AlloyDBOmni(db_config)
     raise ValueError("DB Type not Supported")

evalbench/databases/__init__.py

@@ -16,32 +16,19 @@ def __init__(self, db_config):
         """
         super().__init__(db_config)
         self.nl_config = db_config['nl_config']
-        self.omni = db_config.get("omni", False)
-        self.host = db_config.get("host", "localhost")
-        self.port = db_config.get("port", 5432)
 
         self.connector = AlloyDBConnector()
 
         def get_conn_alloydb():
-            if not self.omni:
-                return self.connector.connect(
-                    self.db_path,
-                    "pg8000",
-                    user=self.username,
-                    password=self.password,
-                    db=self.db_name,
-                    enable_iam_auth=False,
-                    ip_type=AlloyDBIPTypes.PUBLIC,
-                )
-            else:
-                conn_str = f"user={self.username} password={self.password} host={self.host} port={self.port} dbname={self.db_name}"
-                return pg8000.connect(
-                    user=self.username,
-                    password=self.password,
-                    host=self.host,
-                    port=self.port,
-                    database=self.db_name
-                )
+            return self.connector.connect(
+                self.db_path,
+                "pg8000",
+                user=self.username,
+                password=self.password,
+                db=self.db_name,
+                enable_iam_auth=False,
+                ip_type=AlloyDBIPTypes.PUBLIC,
+            )
 
         def get_engine_args_alloydb():
             common_args = {

evalbench/databases/alloydb.py

@@ -0,0 +1,43 @@
+
+from .db import DB
+from .postgres import PGDB
+import sqlalchemy
+import pg8000
+from sqlalchemy.pool import NullPool
+
+
+class AlloyDBOmni(PGDB):
+    def __init__(self, db_config):
+        """
+        Initializes the AlloyDB connection, overriding the PGDB's
+        default Google Cloud SQL connection mechanism.
+        """
+        super().__init__(db_config)
+        self.nl_config = db_config['nl_config']
+        self.host = db_config.get("host", "localhost")
+        self.port = db_config.get("port", 5432)
+
+        def get_conn_alloydb():
+            return pg8000.connect(
+                user=self.username,
+                password=self.password,
+                host=self.host,
+                port=self.port,
+                database=self.db_name
+            )
+
+        def get_engine_args_alloydb():
+            common_args = {
+                "creator": get_conn_alloydb,
+                "connect_args": {"command_timeout": 60},
+            }
+            if "is_tmp_db" in db_config:
+                common_args["poolclass"] = NullPool
+            else:
+                common_args["pool_size"] = 50
+                common_args["pool_recycle"] = 300
+            return common_args
+
+        self.engine = sqlalchemy.create_engine(
+            "postgresql+pg8000://", **get_engine_args_alloydb()
+        )