"""seed overview demo data Revision ID: 20260116_03 Revises: 20260116_02 Create Date: 2026-01-16 16:30:00.000000 """ from typing import Sequence, Union from datetime import date, timedelta from alembic import op import sqlalchemy as sa from sqlalchemy.dialects.postgresql import UUID # revision identifiers, used by Alembic. revision: str = '20260116_03' down_revision: Union[str, None] = '20260116_02' branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: # 获取connection conn = op.get_bind() # 首先获取第一个项目ID result = conn.execute(sa.text("SELECT id FROM studies LIMIT 1")) study_row = result.first() if not study_row: print("没有项目,跳过示例数据注入") return study_id = study_row[0] # 获取该项目的中心 result = conn.execute(sa.text("SELECT id, name FROM sites WHERE study_id = :study_id ORDER BY created_at LIMIT 5"), {"study_id": study_id}) sites = result.fetchall() if not sites: print("没有中心,跳过示例数据注入") return print(f"为 {len(sites)} 个中心注入示例数据...") # 清理可能存在的演示数据(使用ON CONFLICT处理或先删除) # 收集site_ids site_ids = [site[0] for site in sites] # 删除这些中心的existing数据 (如果是演示数据) # 由于我们无法判断哪些是演示数据,使用ON CONFLICT DO NOTHING策略 # 更新入组目标 targets = [80, 60, 70, 50, 40] for idx, site in enumerate(sites): site_id = site[0] target = targets[idx] if idx < len(targets) else 50 conn.execute( sa.text("UPDATE sites SET enrollment_target = :target WHERE id = :site_id"), {"target": target, "site_id": site_id} ) # 基准日期 base_date = date(2024, 12, 1) # 为每个中心添加启动流程数据 for idx, site in enumerate(sites): site_id = site[0] # 机构立项 - 使用 ON CONFLICT DO NOTHING (假设有唯一约束) submit_date = base_date + timedelta(days=idx * 10) accept_date = base_date + timedelta(days=idx * 10 + 5) approved_date = base_date + timedelta(days=idx * 10 + 15) if idx < 4 else None # 检查是否已存在 existing = conn.execute(sa.text( "SELECT 1 FROM startup_feasibility WHERE study_id = :study_id AND site_id = :site_id LIMIT 1" ), {"study_id": study_id, "site_id": site_id}).first() if not existing: conn.execute(sa.text(""" INSERT INTO startup_feasibility (id, study_id, site_id, submit_date, accept_date, approved_date, project_no, created_at) VALUES (gen_random_uuid(), :study_id, :site_id, :submit_date, :accept_date, :approved_date, :project_no, NOW()) """), { "study_id": study_id, "site_id": site_id, "submit_date": submit_date, "accept_date": accept_date, "approved_date": approved_date, "project_no": f"PROJ-{idx+1:03d}" }) # 伦理审批 ethics_submit = base_date + timedelta(days=idx * 10 + 20) ethics_accept = base_date + timedelta(days=idx * 10 + 25) meeting_date = base_date + timedelta(days=idx * 10 + 35) if idx < 4 else None ethics_approved = base_date + timedelta(days=idx * 10 + 40) if idx < 4 else None existing = conn.execute(sa.text( "SELECT 1 FROM startup_ethics WHERE study_id = :study_id AND site_id = :site_id LIMIT 1" ), {"study_id": study_id, "site_id": site_id}).first() if not existing: conn.execute(sa.text(""" INSERT INTO startup_ethics (id, study_id, site_id, submit_date, accept_date, meeting_date, approved_date, approval_no, created_at) VALUES (gen_random_uuid(), :study_id, :site_id, :submit_date, :accept_date, :meeting_date, :approved_date, :approval_no, NOW()) """), { "study_id": study_id, "site_id": site_id, "submit_date": ethics_submit, "accept_date": ethics_accept, "meeting_date": meeting_date, "approved_date": ethics_approved, "approval_no": f"EC-{idx+1:03d}" if idx < 4 else None }) # 合同签署 (前3个中心) if idx < 3: site_name = site[1] signed_date = base_date + timedelta(days=idx * 10 + 50) # 检查合同是否已存在 existing = conn.execute(sa.text( "SELECT 1 FROM finance_contracts WHERE study_id = :study_id AND site_name = :site_name LIMIT 1" ), {"study_id": study_id, "site_name": site_name}).first() if not existing: conn.execute(sa.text(""" INSERT INTO finance_contracts (id, study_id, site_name, contract_no, signed_date, amount, currency, created_at) VALUES (gen_random_uuid(), :study_id, :site_name, :contract_no, :signed_date, 500000.0, 'CNY', NOW()) """), { "study_id": study_id, "site_name": site_name, "contract_no": f"CT-{idx+1:03d}", "signed_date": signed_date }) # 启动会 (前3个中心) if idx < 3: kickoff_date = base_date + timedelta(days=idx * 10 + 60) # 检查启动会是否已存在 existing = conn.execute(sa.text( "SELECT 1 FROM kickoff_meetings WHERE study_id = :study_id AND site_id = :site_id LIMIT 1" ), {"study_id": study_id, "site_id": site_id}).first() if not existing: conn.execute(sa.text(""" INSERT INTO kickoff_meetings (id, study_id, site_id, kickoff_date, attendees, created_at) VALUES (gen_random_uuid(), :study_id, :site_id, :kickoff_date, '["研究者", "CRA", "PM"]'::jsonb, NOW()) """), { "study_id": study_id, "site_id": site_id, "kickoff_date": kickoff_date }) # 为前3个中心添加受试者 enrollment_counts = [52, 18, 70] start_enrollment = date(2025, 1, 1) for idx in range(min(3, len(sites))): site_id = sites[idx][0] count = enrollment_counts[idx] if idx < len(enrollment_counts) else 0 # 检查该中心是否已有受试者数据 existing_count = conn.execute(sa.text( "SELECT COUNT(*) FROM subjects WHERE study_id = :study_id AND site_id = :site_id" ), {"study_id": study_id, "site_id": site_id}).scalar() # 如果已有数据,跳过该中心 if existing_count > 0: print(f" 中心 {idx+1} 已有 {existing_count} 名受试者,跳过") continue for i in range(count): enrollment_date = start_enrollment + timedelta(days=i * 4) screening_date = enrollment_date - timedelta(days=7) consent_date = enrollment_date - timedelta(days=3) conn.execute(sa.text(""" INSERT INTO subjects (id, study_id, site_id, subject_no, status, screening_date, consent_date, enrollment_date, created_at) VALUES (gen_random_uuid(), :study_id, :site_id, :subject_no, 'ENROLLED', :screening_date, :consent_date, :enrollment_date, NOW()) """), { "study_id": study_id, "site_id": site_id, "subject_no": f"S{idx+1:02d}-{i+1:03d}", "screening_date": screening_date, "consent_date": consent_date, "enrollment_date": enrollment_date }) print("✅ 示例数据注入完成") def downgrade() -> None: # 删除示例数据 (可选 - 也可以留空) pass