71 lines
2 KiB
Python
71 lines
2 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Initialize or upgrade the main database for PyraFiles.
|
|
|
|
Usage:
|
|
python initialize_db.py
|
|
|
|
Environment Variables:
|
|
PYRAFILES_DB_URL (optional):
|
|
The SQLAlchemy database URL (e.g., "sqlite:///data/main.db").
|
|
Defaults to "sqlite:///data/main.db" if not set.
|
|
|
|
Description:
|
|
Creates or updates all tables referenced by `Base.metadata`.
|
|
If the database file/tables do not exist, they will be created.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
import transaction
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
# Import Base and DATA_DIR from your app.py
|
|
from app import Base, DATA_DIR # Adjust the import if needed
|
|
|
|
def usage():
|
|
script = os.path.basename(sys.argv[0])
|
|
print(f"Usage: {script}")
|
|
print("Example:")
|
|
print(f" python {script}")
|
|
sys.exit(1)
|
|
|
|
def main():
|
|
logging.basicConfig(level=logging.INFO)
|
|
if len(sys.argv) > 1:
|
|
# We do not expect any arguments
|
|
usage()
|
|
|
|
# Ensure DATA_DIR exists
|
|
if not os.path.exists(DATA_DIR):
|
|
os.makedirs(DATA_DIR)
|
|
logging.info(f"Created data directory at {DATA_DIR}")
|
|
|
|
# Set default database URL to use the data directory
|
|
default_db_url = f"sqlite:///{os.path.join(DATA_DIR, 'main.db')}"
|
|
# Read environment variable for DB URL or use default
|
|
db_url = os.environ.get("PYRAFILES_DB_URL", default_db_url)
|
|
logging.info(f"Using DB URL: {db_url}")
|
|
|
|
# Set up engine
|
|
engine = create_engine(
|
|
db_url,
|
|
connect_args={"check_same_thread": False} if "sqlite" in db_url else {},
|
|
poolclass=StaticPool if "sqlite" in db_url else None,
|
|
)
|
|
|
|
SessionFactory = sessionmaker(bind=engine)
|
|
|
|
# Create or upgrade tables
|
|
with transaction.manager:
|
|
logging.info("Creating or upgrading tables using Base.metadata.create_all()")
|
|
Base.metadata.create_all(engine)
|
|
|
|
logging.info("Database initialization complete.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|