68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
"""add experiment project_id column
|
||
|
||
Revision ID: add_experiment_project_id
|
||
Revises: add_budget_tables
|
||
Create Date: 2026-02-25
|
||
"""
|
||
|
||
from alembic import op
|
||
import sqlalchemy as sa
|
||
|
||
revision = "add_experiment_project_id"
|
||
down_revision = "add_budget_tables"
|
||
branch_labels = None
|
||
depends_on = None
|
||
|
||
|
||
def _has_column(table_name: str, column_name: str) -> bool:
|
||
bind = op.get_bind()
|
||
inspector = sa.inspect(bind)
|
||
columns = inspector.get_columns(table_name)
|
||
return any(col["name"] == column_name for col in columns)
|
||
|
||
|
||
def _has_index(table_name: str, index_name: str) -> bool:
|
||
bind = op.get_bind()
|
||
inspector = sa.inspect(bind)
|
||
indexes = inspector.get_indexes(table_name)
|
||
return any(idx["name"] == index_name for idx in indexes)
|
||
|
||
|
||
def upgrade():
|
||
if not _has_column("experiments", "project_id"):
|
||
op.add_column("experiments", sa.Column("project_id", sa.BigInteger(), nullable=True))
|
||
|
||
bind = op.get_bind()
|
||
dialect = bind.dialect.name
|
||
|
||
# 尽量从 linked_context_nodes 推断 project_id（仅唯一归属时回填）
|
||
if dialect == "sqlite":
|
||
op.execute(
|
||
"""
|
||
UPDATE experiments
|
||
SET project_id = (
|
||
SELECT tn.project_id
|
||
FROM truth_nodes tn
|
||
WHERE instr(COALESCE(experiments.linked_context_nodes, ''), printf('%s', tn.id)) > 0
|
||
AND tn.project_id IS NOT NULL
|
||
GROUP BY tn.project_id
|
||
HAVING COUNT(*) >= 1
|
||
LIMIT 1
|
||
)
|
||
WHERE project_id IS NULL
|
||
"""
|
||
)
|
||
else:
|
||
# PostgreSQL JSON/数组形态不稳定，迁移阶段不做激进回填，留给服务层按需补写。
|
||
pass
|
||
|
||
if not _has_index("experiments", "ix_experiments_project_id"):
|
||
op.create_index("ix_experiments_project_id", "experiments", ["project_id"])
|
||
|
||
|
||
def downgrade():
|
||
if _has_index("experiments", "ix_experiments_project_id"):
|
||
op.drop_index("ix_experiments_project_id", table_name="experiments")
|
||
if _has_column("experiments", "project_id"):
|
||
op.drop_column("experiments", "project_id")
|