68 lines
2.6 KiB
Python
68 lines
2.6 KiB
Python
# database.py
|
||
import os
|
||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
|
||
from sqlalchemy.orm import sessionmaker, declarative_base
|
||
|
||
# 优先从环境变量读取（Docker 部署推荐）
|
||
SQLALCHEMY_DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+asyncpg://user:password@db:5432/fastapi_db")
|
||
|
||
engine = create_async_engine(SQLALCHEMY_DATABASE_URL, echo=True)
|
||
AsyncSessionLocal = sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False)
|
||
|
||
Base = declarative_base()
|
||
|
||
async def get_db():
|
||
async with AsyncSessionLocal() as session:
|
||
yield session
|
||
# 注意：通常不在 get_db 里统一 commit，建议在 endpoint 里手动 commit
|
||
|
||
async def init_db():
|
||
async with engine.begin() as conn:
|
||
# 导入模型以确保 metadata 注册了表
|
||
from . import models
|
||
# 自动创建表
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
|
||
# 创建默认管理员用户
|
||
async with AsyncSessionLocal() as session:
|
||
from . import auth
|
||
from sqlalchemy.future import select
|
||
|
||
# 检查是否已存在 admin 用户
|
||
result = await session.execute(
|
||
select(models.User).where(models.User.username == "admin")
|
||
)
|
||
existing_admin = result.scalars().first()
|
||
|
||
if not existing_admin:
|
||
# 创建管理员用户
|
||
private_key, public_key = auth.generate_key_pair()
|
||
admin_user = models.User(
|
||
username="admin",
|
||
hashed_password=auth.hash_password("admin123"),
|
||
private_key=private_key,
|
||
public_key=public_key,
|
||
is_admin=True,
|
||
guale=False
|
||
)
|
||
session.add(admin_user)
|
||
await session.commit()
|
||
print("✅ Default admin user created (username: admin, password: admin123)")
|
||
|
||
# 检查是否已存在 Gemini 配置
|
||
result = await session.execute(
|
||
select(models.AIConfig).where(models.AIConfig.provider_name == "gemini")
|
||
)
|
||
existing_gemini = result.scalars().first()
|
||
|
||
if not existing_gemini:
|
||
gemini_config = models.AIConfig(
|
||
provider_name="gemini",
|
||
api_key=os.getenv("GEMINI_API_KEY", "your-gemini-api-key"),
|
||
api_url="https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
|
||
default_model="gemini-3-flash-preview",
|
||
is_active=True
|
||
)
|
||
session.add(gemini_config)
|
||
await session.commit()
|
||
print("✅ Default Gemini AI configuration created") |