142 lines
No EOL
5 KiB
Python
142 lines
No EOL
5 KiB
Python
import os
|
|
import logging
|
|
import psycopg2
|
|
from dotenv import load_dotenv
|
|
from relationaldb.database import Base
|
|
from sqlalchemy import create_engine, text
|
|
|
|
from relationaldb.models import memory
|
|
from relationaldb.models import metadatas
|
|
from relationaldb.models import operation
|
|
from relationaldb.models import sessions
|
|
from relationaldb.models import user
|
|
from relationaldb.models import docs
|
|
|
|
load_dotenv()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
import os
|
|
import logging
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from contextlib import contextmanager
|
|
from dotenv import load_dotenv
|
|
from relationaldb.database import Base # Assuming all models are imported within this module
|
|
from relationaldb.database import DatabaseConfig # Assuming DatabaseConfig is defined as before
|
|
|
|
load_dotenv()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class DatabaseManager:
|
|
def __init__(self, config: DatabaseConfig):
|
|
self.config = config
|
|
self.engine = create_engine(config.get_sqlalchemy_database_url())
|
|
self.db_type = config.db_type
|
|
|
|
@contextmanager
|
|
def get_connection(self):
|
|
if self.db_type in ['sqlite', 'duckdb']:
|
|
# For SQLite and DuckDB, the engine itself manages connections
|
|
yield self.engine
|
|
else:
|
|
connection = self.engine.connect()
|
|
try:
|
|
yield connection
|
|
finally:
|
|
connection.close()
|
|
|
|
def database_exists(self, db_name):
|
|
if self.db_type in ['sqlite', 'duckdb']:
|
|
# For SQLite and DuckDB, check if the database file exists
|
|
return os.path.exists(db_name)
|
|
else:
|
|
query = text(f"SELECT 1 FROM pg_database WHERE datname='{db_name}'")
|
|
with self.get_connection() as connection:
|
|
result = connection.execute(query).fetchone()
|
|
return result is not None
|
|
|
|
def create_database(self, db_name):
|
|
if self.db_type not in ['sqlite', 'duckdb']:
|
|
# For databases like PostgreSQL, create the database explicitly
|
|
with self.get_connection() as connection:
|
|
connection.execution_options(isolation_level="AUTOCOMMIT")
|
|
connection.execute(f"CREATE DATABASE {db_name}")
|
|
|
|
def drop_database(self, db_name):
|
|
if self.db_type in ['sqlite', 'duckdb']:
|
|
# For SQLite and DuckDB, simply remove the database file
|
|
os.remove(db_name)
|
|
else:
|
|
with self.get_connection() as connection:
|
|
connection.execution_options(isolation_level="AUTOCOMMIT")
|
|
connection.execute(f"DROP DATABASE IF EXISTS {db_name}")
|
|
|
|
def create_tables(self):
|
|
Base.metadata.create_all(bind=self.engine)
|
|
|
|
if __name__ == "__main__":
|
|
# Example usage with SQLite
|
|
config = DatabaseConfig(db_type='sqlite', db_name='mydatabase.db')
|
|
|
|
# For DuckDB, you would set db_type to 'duckdb' and provide the database file name
|
|
# config = DatabaseConfig(db_type='duckdb', db_name='mydatabase.duckdb')
|
|
|
|
db_manager = DatabaseManager(config=config)
|
|
|
|
database_name = config.db_name
|
|
|
|
if not db_manager.database_exists(database_name):
|
|
logger.info(f"Database {database_name} does not exist. Creating...")
|
|
db_manager.create_database(database_name)
|
|
logger.info(f"Database {database_name} created successfully.")
|
|
|
|
db_manager.create_tables()
|
|
|
|
#
|
|
# def create_admin_engine(username, password, host, database_name):
|
|
# admin_url = f"postgresql://{username}:{password}@{host}:5432/{database_name}"
|
|
# return create_engine(admin_url)
|
|
#
|
|
# def database_exists(connection, db_name):
|
|
# query = text(f"SELECT 1 FROM pg_database WHERE datname='{db_name}'")
|
|
# result = connection.execute(query).fetchone()
|
|
# return result is not None
|
|
#
|
|
# def create_database(connection, db_name):
|
|
# connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
|
|
# cursor = connection.cursor()
|
|
# cursor.execute(f"CREATE DATABASE {db_name}")
|
|
# cursor.close()
|
|
#
|
|
# def drop_database(connection, db_name):
|
|
# connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
|
|
# cursor = connection.cursor()
|
|
# cursor.execute(f"DROP DATABASE IF EXISTS {db_name}")
|
|
# cursor.close()
|
|
#
|
|
#
|
|
#
|
|
# def create_tables(engine):
|
|
# Base.metadata.create_all(bind = engine)
|
|
#
|
|
# if __name__ == "__main__":
|
|
# host = os.environ.get('POSTGRES_HOST')
|
|
# username = os.environ.get('POSTGRES_USER')
|
|
# password = os.environ.get('POSTGRES_PASSWORD')
|
|
# database_name = os.environ.get('POSTGRES_DB')
|
|
#
|
|
# engine = create_admin_engine(username, password, host, database_name)
|
|
# connection = engine.connect()
|
|
#
|
|
# # print(Base.metadata.tables)
|
|
#
|
|
# if not database_exists(connection, database_name):
|
|
# logger.info(f"Database {database_name} does not exist. Creating...")
|
|
# create_database(connection, database_name)
|
|
# logger.info(f"Database {database_name} created successfully.")
|
|
#
|
|
# connection.close()
|
|
# engine.dispose()
|
|
#
|
|
# create_tables(engine) |