diff options
Diffstat (limited to 'database.py')
| -rw-r--r-- | database.py | 17 |
1 files changed, 11 insertions, 6 deletions
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() |