From 4d626d423c2e8158205fb525f24f3dc2338a5b1f Mon Sep 17 00:00:00 2001 From: Parker Date: Sat, 16 Nov 2024 01:48:58 -0600 Subject: Add support for MySQL and PostgreSQL --- database.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) (limited to 'database.py') diff --git a/database.py b/database.py index 544ee05..0166d28 100644 --- a/database.py +++ b/database.py @@ -1,13 +1,18 @@ from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker -import os -# Create 'data' directory at root if it doesn't exist -if not os.path.exists("data"): - os.makedirs("data") +import config -engine = create_engine("sqlite:///data/data.db") -SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) +if config.DB_ENGINE == "mysql": + database_url = f"mysql+pymysql://{config.DB_USER}:{config.DB_PASSWORD}@{config.DB_HOST}:{config.DB_PORT}/{config.DB_NAME}" + +elif config.DB_ENGINE == "postgresql": + database_url = f"postgresql+psycopg2://{config.DB_USER}:{config.DB_PASSWORD}@{config.DB_HOST}:{config.DB_PORT}/{config.DB_NAME}" +else: + database_url = "sqlite:///data/data.db" + +engine = create_engine(database_url) +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() -- cgit v1.2.3-70-g09d2