115 lines
4.3 KiB
Python
115 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Delete expired project workbench events in bounded batches."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
import sys
|
|
from uuid import UUID
|
|
|
|
from psycopg import AsyncConnection
|
|
from psycopg.conninfo import conninfo_to_dict
|
|
|
|
REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(REPOSITORY_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPOSITORY_ROOT))
|
|
|
|
from app.auth.project_dependencies import DB_ROLE_BIZ_DATA
|
|
from app.core.config import settings
|
|
from app.infra.db.metadb.database import SessionLocal, close_metadata_engine
|
|
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
|
|
|
|
|
|
def _database_identity(dsn: str) -> tuple[str, str, str]:
|
|
values = conninfo_to_dict(dsn)
|
|
return (
|
|
values.get("host") or values.get("hostaddr") or "localhost",
|
|
values.get("port") or "5432",
|
|
values.get("dbname") or values.get("user") or "",
|
|
)
|
|
|
|
|
|
async def cleanup_project(
|
|
project_id: UUID, dsn: str, retention_days: int, batch_size: int, dry_run: bool
|
|
) -> int:
|
|
async with await AsyncConnection.connect(dsn) as conn:
|
|
if dry_run:
|
|
cursor = await conn.execute(
|
|
"""SELECT count(*) FROM workbench.domain_event_log
|
|
WHERE project_id=%s
|
|
AND occurred_at < now() - make_interval(days => %s)""",
|
|
(project_id, retention_days),
|
|
)
|
|
return int((await cursor.fetchone())[0])
|
|
deleted = 0
|
|
while True:
|
|
cursor = await conn.execute(
|
|
"""WITH expired AS (
|
|
SELECT event_seq FROM workbench.domain_event_log
|
|
WHERE project_id=%s
|
|
AND occurred_at < now() - make_interval(days => %s)
|
|
ORDER BY event_seq LIMIT %s
|
|
)
|
|
DELETE FROM workbench.domain_event_log e
|
|
USING expired WHERE e.event_seq=expired.event_seq""",
|
|
(project_id, retention_days, batch_size),
|
|
)
|
|
await conn.commit()
|
|
count = cursor.rowcount
|
|
deleted += count
|
|
if count < batch_size:
|
|
return deleted
|
|
|
|
|
|
async def cleanup(
|
|
retention_days: int, batch_size: int, project_id: UUID | None, dry_run: bool
|
|
) -> dict[UUID, int]:
|
|
async with SessionLocal() as session:
|
|
repository = MetadataRepository(session)
|
|
routings = await repository.list_active_project_db_routings(DB_ROLE_BIZ_DATA)
|
|
if project_id is not None:
|
|
routings = [routing for routing in routings if routing.project_id == project_id]
|
|
if not routings:
|
|
raise ValueError("active project business database routing was not found")
|
|
by_database = defaultdict(list)
|
|
for routing in routings:
|
|
by_database[_database_identity(routing.dsn)].append(routing)
|
|
shared = [items for items in by_database.values() if len(items) > 1]
|
|
if shared:
|
|
project_ids = ", ".join(str(item.project_id) for items in shared for item in items)
|
|
raise ValueError(f"shared business database routing is unsupported: {project_ids}")
|
|
return {
|
|
routing.project_id: await cleanup_project(
|
|
routing.project_id, routing.dsn, retention_days, batch_size, dry_run
|
|
)
|
|
for routing in routings
|
|
}
|
|
|
|
|
|
async def main(args) -> None:
|
|
try:
|
|
results = await cleanup(
|
|
args.retention_days, args.batch_size, args.project_id, args.dry_run
|
|
)
|
|
for project_id, count in results.items():
|
|
action = "would delete" if args.dry_run else "deleted"
|
|
print(f"{project_id}: {action} {count} events")
|
|
finally:
|
|
await close_metadata_engine()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--retention-days", type=int, default=settings.WORKBENCH_EVENT_RETENTION_DAYS
|
|
)
|
|
parser.add_argument("--batch-size", type=int, default=5000)
|
|
parser.add_argument("--project-id", type=UUID)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.retention_days < 1 or args.batch_size < 1:
|
|
parser.error("retention days and batch size must be positive")
|
|
asyncio.run(main(args))
|