|
2 | 2 |
|
3 | 3 | import logging |
4 | 4 | from importlib import import_module |
5 | | -from typing import Tuple |
6 | 5 |
|
7 | 6 | from sqlalchemy import create_engine, inspect, text |
8 | 7 | from sqlalchemy.orm import declarative_base, sessionmaker |
|
21 | 20 | SESSION_LOCAL = sessionmaker(autocommit=False, autoflush=False, bind=engine) |
22 | 21 |
|
23 | 22 |
|
24 | | -def _timestamp_sql(dialect_name: str) -> Tuple[str, str]: |
25 | | - """Return dialect-specific SQL fragments for timestamp columns.""" |
26 | | - |
27 | | - if dialect_name == "postgresql": |
28 | | - return "TIMESTAMPTZ", "NOW()" |
29 | | - if dialect_name == "sqlite": |
30 | | - return "TIMESTAMP", "CURRENT_TIMESTAMP" |
31 | | - return "TIMESTAMP", "CURRENT_TIMESTAMP" |
32 | | - |
33 | | - |
34 | | -def _migrate_legacy_users_table() -> None: |
35 | | - """Rename and backfill the legacy ``users`` table if it exists.""" |
36 | | - |
37 | | - with engine.begin() as connection: |
38 | | - inspector = inspect(connection) |
39 | | - existing_tables = set(inspector.get_table_names()) |
40 | | - if "user" in existing_tables or "users" not in existing_tables: |
41 | | - return |
42 | | - |
43 | | - logger.info("🚚 Migrating legacy users table to new schema...") |
44 | | - connection.execute(text('ALTER TABLE "users" RENAME TO "user"')) |
45 | | - |
46 | | - inspector = inspect(connection) |
47 | | - columns = {column["name"] for column in inspector.get_columns("user")} |
48 | | - |
49 | | - if "id" in columns: |
50 | | - connection.execute(text('ALTER TABLE "user" RENAME COLUMN id TO user_id')) |
51 | | - |
52 | | - if "created_at" not in columns: |
53 | | - column_type, default_expr = _timestamp_sql(connection.dialect.name) |
54 | | - connection.execute( |
55 | | - text( |
56 | | - f'ALTER TABLE "user" ADD COLUMN created_at {column_type} DEFAULT {default_expr}' |
57 | | - ) |
58 | | - ) |
59 | | - connection.execute( |
60 | | - text(f'UPDATE "user" SET created_at = {default_expr} WHERE created_at IS NULL') |
61 | | - ) |
62 | | - |
63 | | - logger.info("✅ Legacy users table migrated.") |
64 | | - |
65 | | - |
66 | 23 | def init_db() -> None: |
67 | 24 | """Automatically create or update tables based on SQLAlchemy models.""" |
68 | 25 |
|
69 | | - _migrate_legacy_users_table() |
70 | | - |
71 | 26 | for module in ( |
72 | 27 | "app.models.user", |
73 | 28 | "app.models.project", |
|
0 commit comments