64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import text
|
|
|
|
from app.core.config import settings
|
|
from app.db.session import Base, engine
|
|
from app.models import Asset, Job, JobEvent, JobOutput, User # noqa: F401
|
|
from app.services.storage import asset_storage
|
|
|
|
|
|
def init_db() -> None:
|
|
Path(settings.ASSET_STORAGE_ROOT).mkdir(parents=True, exist_ok=True)
|
|
Path(settings.OUTPUT_STORAGE_ROOT).mkdir(parents=True, exist_ok=True)
|
|
Base.metadata.create_all(bind=engine)
|
|
_migrate_assets_table()
|
|
_cleanup_expired_trashed_assets()
|
|
|
|
|
|
def _migrate_assets_table() -> None:
|
|
with engine.begin() as conn:
|
|
columns = {
|
|
row[1]
|
|
for row in conn.execute(text("PRAGMA table_info(assets)")).fetchall()
|
|
}
|
|
if "is_trashed" not in columns:
|
|
conn.execute(text("ALTER TABLE assets ADD COLUMN is_trashed BOOLEAN NOT NULL DEFAULT 0"))
|
|
if "delete_after_at" not in columns:
|
|
conn.execute(text("ALTER TABLE assets ADD COLUMN delete_after_at DATETIME NULL"))
|
|
|
|
|
|
def _cleanup_expired_trashed_assets() -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
with engine.begin() as conn:
|
|
rows = conn.execute(
|
|
text(
|
|
"""
|
|
SELECT id, storage_path, thumbnail_path
|
|
FROM assets
|
|
WHERE is_trashed = 1
|
|
AND delete_after_at IS NOT NULL
|
|
AND delete_after_at <= :now
|
|
"""
|
|
),
|
|
{"now": now},
|
|
).fetchall()
|
|
|
|
for _, storage_path, thumbnail_path in rows:
|
|
asset_storage.delete_relative_path(storage_path)
|
|
asset_storage.delete_relative_path(thumbnail_path)
|
|
|
|
if rows:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
DELETE FROM assets
|
|
WHERE is_trashed = 1
|
|
AND delete_after_at IS NOT NULL
|
|
AND delete_after_at <= :now
|
|
"""
|
|
),
|
|
{"now": now},
|
|
)
|