Add SQL database creation function

bnpm/sql_helpers.py

@@ -1,9 +1,9 @@
 from typing import Optional, Union, List, Tuple, Dict, Any, Iterable, Sequence
 
+import urllib
+
 import sqlalchemy
-import pymysql
 import pandas as pd
-import urllib
 
 
 def format_url(
@@ -337,4 +337,33 @@ def get_table_data(
     if limit is not None:
         query += f" LIMIT {limit}"
 
-    return pd.read_sql(query, connection)
+    return pd.read_sql(query, connection)
+
+
+def create_database(
+    connection: sqlalchemy.engine.base.Connection,
+    database: str,
+) -> None:
+    """
+    Creates a new database on the server.
+    RH 2024
+    
+    Args:
+        connection (sqlalchemy.engine.base.Connection):
+            SQL connection object
+        database (str):
+            Name of the database to create
+    """
+    if isinstance(connection, sqlalchemy.engine.base.Connection):
+        if "mysql" in str(connection.engine):
+            query = f"CREATE DATABASE {database}"
+        elif "postgresql" in str(connection.engine):
+            query = f"CREATE DATABASE {database}"
+        elif "sqlite" in str(connection.engine):
+            raise ValueError("SQLite does not support this operation")
+        else:
+            raise ValueError("Connection type not recognized")
+    else:
+        raise TypeError("connection must be a sqlalchemy.engine.base.Connection object")
+
+    connection.execute(query)

requirements.txt

@@ -20,6 +20,7 @@ tqdm
 seaborn
 sparse
 scipy
+sqlalchemy
 xxhash
 wandb
 torch

setup.py

@@ -80,6 +80,7 @@ def read_requirements():
     'ipywidgets',
     'eva_decord',
     'wandb',
+    'sqlalchemy',
 ]}
 
 deps_advanced = {dep: deps_all_dict[dep] for dep in [