83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
"""Add user_department association table migration."""
|
|
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# 添加项目根目录到Python路径
|
|
project_root = Path(__file__).parent.parent.parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
import asyncio
|
|
import asyncpg
|
|
from th_agenter.core.config import get_settings
|
|
|
|
async def create_user_department_table():
|
|
"""Create user_departments association table."""
|
|
settings = get_settings()
|
|
database_url = settings.database.url
|
|
|
|
print(f"Database URL: {database_url}")
|
|
|
|
try:
|
|
# 解析PostgreSQL连接URL
|
|
# postgresql://user:password@host:port/database
|
|
url_parts = database_url.replace('postgresql://', '').split('/')
|
|
db_name = url_parts[1] if len(url_parts) > 1 else 'postgres'
|
|
user_host = url_parts[0].split('@')
|
|
user_pass = user_host[0].split(':')
|
|
host_port = user_host[1].split(':')
|
|
|
|
user = user_pass[0]
|
|
password = user_pass[1] if len(user_pass) > 1 else ''
|
|
host = host_port[0]
|
|
port = int(host_port[1]) if len(host_port) > 1 else 5432
|
|
|
|
# 连接PostgreSQL数据库
|
|
conn = await asyncpg.connect(
|
|
user=user,
|
|
password=password,
|
|
database=db_name,
|
|
host=host,
|
|
port=port
|
|
)
|
|
|
|
# 创建user_departments表
|
|
create_table_sql = """
|
|
CREATE TABLE IF NOT EXISTS user_departments (
|
|
id SERIAL PRIMARY KEY,
|
|
user_id INTEGER NOT NULL,
|
|
department_id INTEGER NOT NULL,
|
|
is_primary BOOLEAN NOT NULL DEFAULT true,
|
|
is_active BOOLEAN NOT NULL DEFAULT true,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
|
|
FOREIGN KEY (department_id) REFERENCES departments (id) ON DELETE CASCADE
|
|
);
|
|
"""
|
|
|
|
await conn.execute(create_table_sql)
|
|
|
|
# 创建索引
|
|
create_indexes_sql = [
|
|
"CREATE INDEX IF NOT EXISTS idx_user_departments_user_id ON user_departments (user_id);",
|
|
"CREATE INDEX IF NOT EXISTS idx_user_departments_department_id ON user_departments (department_id);",
|
|
"CREATE UNIQUE INDEX IF NOT EXISTS idx_user_departments_unique ON user_departments (user_id, department_id);"
|
|
]
|
|
|
|
for index_sql in create_indexes_sql:
|
|
await conn.execute(index_sql)
|
|
|
|
print("User departments table created successfully")
|
|
|
|
except Exception as e:
|
|
print(f"Error creating user departments table: {e}")
|
|
raise
|
|
finally:
|
|
if 'conn' in locals():
|
|
await conn.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(create_user_department_table()) |