i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: 修复迁移脚本

Browse Source
This commit is contained in:
fawney19
2026-02-28 14:03:42 +08:00
parent 54530faf03
commit 1d644de500
1 changed files with 58 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
alembic/versions/20260228_1100_b2c3d4e5f6a7_cleanup_capability_settings.py
View File

@@ -33,11 +33,23 @@ depends_on: str | Sequence[str] | None = None
def column_exists(table_name: str, column_name: str) -> bool: def column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind() bind = op.get_bind()
inspector = inspect(bind) insp = inspect(bind)
columns = [c["name"] for c in inspector.get_columns(table_name)] columns = [c["name"] for c in insp.get_columns(table_name)]
return column_name in columns return column_name in columns
def table_exists(table_name: str) -> bool:
bind = op.get_bind()
insp = inspect(bind)
return table_name in insp.get_table_names()
def index_exists(table_name: str, index_name: str) -> bool:
bind = op.get_bind()
insp = inspect(bind)
return any(idx["name"] == index_name for idx in insp.get_indexes(table_name))
def upgrade() -> None: def upgrade() -> None:
# --- 1. Add cache_creation columns --- # --- 1. Add cache_creation columns ---
if not column_exists("usage", "cache_creation_input_tokens_5m"): if not column_exists("usage", "cache_creation_input_tokens_5m"):
@@ -126,33 +138,36 @@ def upgrade() -> None:
) )
# --- 3. Create user_model_usage_counts table --- # --- 3. Create user_model_usage_counts table ---
op.create_table( if not table_exists("user_model_usage_counts"):
"user_model_usage_counts", op.create_table(
sa.Column("id", sa.String(36), primary_key=True), "user_model_usage_counts",
sa.Column( sa.Column("id", sa.String(36), primary_key=True),
"user_id", sa.Column(
sa.String(36), "user_id",
sa.ForeignKey("users.id", ondelete="CASCADE"), sa.String(36),
nullable=False, sa.ForeignKey("users.id", ondelete="CASCADE"),
), nullable=False,
sa.Column("model", sa.String(100), nullable=False), ),
sa.Column("usage_count", sa.Integer, nullable=False, server_default="0"), sa.Column("model", sa.String(100), nullable=False),
sa.Column( sa.Column("usage_count", sa.Integer, nullable=False, server_default="0"),
"created_at", sa.Column(
sa.DateTime(timezone=True), "created_at",
nullable=False, sa.DateTime(timezone=True),
server_default=sa.func.now(), nullable=False,
), server_default=sa.func.now(),
sa.Column( ),
"updated_at", sa.Column(
sa.DateTime(timezone=True), "updated_at",
nullable=False, sa.DateTime(timezone=True),
server_default=sa.func.now(), nullable=False,
), server_default=sa.func.now(),
sa.UniqueConstraint("user_id", "model", name="uq_user_model_usage_count"), ),
) sa.UniqueConstraint("user_id", "model", name="uq_user_model_usage_count"),
op.create_index("idx_user_model_usage_user", "user_model_usage_counts", ["user_id"]) )
op.create_index("idx_user_model_usage_model", "user_model_usage_counts", ["model"]) if not index_exists("user_model_usage_counts", "idx_user_model_usage_user"):
op.create_index("idx_user_model_usage_user", "user_model_usage_counts", ["user_id"])
if not index_exists("user_model_usage_counts", "idx_user_model_usage_model"):
op.create_index("idx_user_model_usage_model", "user_model_usage_counts", ["model"])
# Backfill from existing usage records # Backfill from existing usage records
rows = conn.execute( rows = conn.execute(
@@ -179,18 +194,26 @@ def upgrade() -> None:
) )
# --- 4. Enforce models.global_model_id NOT NULL --- # --- 4. Enforce models.global_model_id NOT NULL ---
op.execute("DELETE FROM models WHERE global_model_id IS NULL") conn = op.get_bind()
op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=False) insp = inspect(conn)
model_cols = {c["name"]: c for c in insp.get_columns("models")}
if model_cols.get("global_model_id", {}).get("nullable", True):
op.execute("DELETE FROM models WHERE global_model_id IS NULL")
op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=False)
def downgrade() -> None: def downgrade() -> None:
# Revert models.global_model_id to nullable # Revert models.global_model_id to nullable
op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=True) if column_exists("models", "global_model_id"):
op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=True)
# Drop user_model_usage_counts # Drop user_model_usage_counts
op.drop_index("idx_user_model_usage_model", table_name="user_model_usage_counts") if table_exists("user_model_usage_counts"):
op.drop_index("idx_user_model_usage_user", table_name="user_model_usage_counts") if index_exists("user_model_usage_counts", "idx_user_model_usage_model"):
op.drop_table("user_model_usage_counts") op.drop_index("idx_user_model_usage_model", table_name="user_model_usage_counts")
if index_exists("user_model_usage_counts", "idx_user_model_usage_user"):
op.drop_index("idx_user_model_usage_user", table_name="user_model_usage_counts")
op.drop_table("user_model_usage_counts")
# Drop cache_creation columns # Drop cache_creation columns
if column_exists("usage", "cache_creation_input_tokens_1h"): if column_exists("usage", "cache_creation_input_tokens_1h"):