feat: video-flow initial commit
- app.py: Streamlit UI for video generation workflow - main_flow.py: CLI tool with argparse support - modules/: Business logic modules (script_gen, image_gen, video_gen, composer, etc.) - config.py: Configuration with API keys and paths - requirements.txt: Python dependencies - docs/: System prompt documentation
This commit is contained in:
Tony Zhang
305
modules/db_manager.py
Normal file
305
modules/db_manager.py
Normal file
@@ -0,0 +1,305 @@
|
||||
"""
|
||||
数据库管理模块 (SQLAlchemy)
|
||||
负责项目数据、任务状态、素材路径的持久化存储
|
||||
支持 SQLite 和 PostgreSQL
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Dict, List, Any, Optional
|
||||
|
||||
from sqlalchemy import create_engine, Column, String, Integer, Text, Float, UniqueConstraint, func
|
||||
from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
class Project(Base):
|
||||
__tablename__ = 'projects'
|
||||
|
||||
id = Column(String, primary_key=True)
|
||||
name = Column(String)
|
||||
status = Column(String) # created, script_generated, images_generated, videos_generated, completed
|
||||
product_info = Column(Text) # JSON string (SQLite) or JSONB (PG - using Text for compat)
|
||||
script_data = Column(Text) # JSON string
|
||||
created_at = Column(Float, default=time.time)
|
||||
updated_at = Column(Float, default=time.time, onupdate=time.time)
|
||||
|
||||
class SceneAsset(Base):
|
||||
__tablename__ = 'scene_assets'
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
project_id = Column(String, index=True)
|
||||
scene_id = Column(Integer)
|
||||
asset_type = Column(String) # image, video
|
||||
status = Column(String) # pending, processing, completed, failed
|
||||
local_path = Column(Text, nullable=True)
|
||||
remote_url = Column(Text, nullable=True)
|
||||
task_id = Column(String, nullable=True) # 外部 API 的任务 ID
|
||||
metadata_json = Column("metadata", Text, nullable=True) # JSON string (renamed to avoid conflict with metadata attr)
|
||||
created_at = Column(Float, default=time.time)
|
||||
updated_at = Column(Float, default=time.time, onupdate=time.time)
|
||||
|
||||
__table_args__ = (UniqueConstraint('project_id', 'scene_id', 'asset_type', name='uix_project_scene_asset'),)
|
||||
|
||||
class AppConfig(Base):
|
||||
__tablename__ = 'app_config'
|
||||
|
||||
key = Column(String, primary_key=True)
|
||||
value = Column(Text) # JSON string
|
||||
description = Column(Text, nullable=True)
|
||||
updated_at = Column(Float, default=time.time, onupdate=time.time)
|
||||
|
||||
class DBManager:
|
||||
def __init__(self, connection_string: str = None):
|
||||
if not connection_string:
|
||||
connection_string = config.DB_CONNECTION_STRING
|
||||
|
||||
self.engine = create_engine(connection_string, pool_recycle=3600)
|
||||
self.Session = scoped_session(sessionmaker(bind=self.engine))
|
||||
self._init_db()
|
||||
|
||||
def _init_db(self):
|
||||
"""初始化表结构"""
|
||||
Base.metadata.create_all(self.engine)
|
||||
|
||||
def _get_session(self):
|
||||
return self.Session()
|
||||
|
||||
# --- Project Operations ---
|
||||
|
||||
def create_project(self, project_id: str, name: str, product_info: Dict[str, Any]):
|
||||
session = self._get_session()
|
||||
try:
|
||||
# Check if exists
|
||||
existing = session.query(Project).filter_by(id=project_id).first()
|
||||
if existing:
|
||||
logger.warning(f"Project {project_id} already exists.")
|
||||
return
|
||||
|
||||
new_project = Project(
|
||||
id=project_id,
|
||||
name=name,
|
||||
status="created",
|
||||
product_info=json.dumps(product_info, ensure_ascii=False),
|
||||
created_at=time.time(),
|
||||
updated_at=time.time()
|
||||
)
|
||||
session.add(new_project)
|
||||
session.commit()
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
logger.error(f"Error creating project: {e}")
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def update_project_script(self, project_id: str, script: Dict[str, Any]):
|
||||
session = self._get_session()
|
||||
try:
|
||||
project = session.query(Project).filter_by(id=project_id).first()
|
||||
if project:
|
||||
project.script_data = json.dumps(script, ensure_ascii=False)
|
||||
project.status = "script_generated"
|
||||
project.updated_at = time.time()
|
||||
session.commit()
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
logger.error(f"Error updating script: {e}")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def update_project_status(self, project_id: str, status: str):
|
||||
session = self._get_session()
|
||||
try:
|
||||
project = session.query(Project).filter_by(id=project_id).first()
|
||||
if project:
|
||||
project.status = status
|
||||
project.updated_at = time.time()
|
||||
session.commit()
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
logger.error(f"Error updating status: {e}")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def get_project(self, project_id: str) -> Optional[Dict[str, Any]]:
|
||||
session = self._get_session()
|
||||
try:
|
||||
project = session.query(Project).filter_by(id=project_id).first()
|
||||
if project:
|
||||
data = {
|
||||
"id": project.id,
|
||||
"name": project.name,
|
||||
"status": project.status,
|
||||
"product_info": json.loads(project.product_info) if project.product_info else {},
|
||||
"script_data": json.loads(project.script_data) if project.script_data else None,
|
||||
"created_at": project.created_at,
|
||||
"updated_at": project.updated_at
|
||||
}
|
||||
return data
|
||||
return None
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def list_projects(self) -> List[Dict[str, Any]]:
|
||||
session = self._get_session()
|
||||
try:
|
||||
projects = session.query(Project).order_by(Project.updated_at.desc()).all()
|
||||
results = []
|
||||
for p in projects:
|
||||
results.append({
|
||||
"id": p.id,
|
||||
"name": p.name,
|
||||
"status": p.status,
|
||||
"updated_at": p.updated_at
|
||||
})
|
||||
return results
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
# --- Asset/Task Operations ---
|
||||
|
||||
def save_asset(self, project_id: str, scene_id: int, asset_type: str,
|
||||
status: str, local_path: str = None, remote_url: str = None,
|
||||
task_id: str = None, metadata: Dict = None):
|
||||
"""保存或更新资产记录 (UPSERT 逻辑)"""
|
||||
session = self._get_session()
|
||||
try:
|
||||
asset = session.query(SceneAsset).filter_by(
|
||||
project_id=project_id,
|
||||
scene_id=scene_id,
|
||||
asset_type=asset_type
|
||||
).first()
|
||||
|
||||
meta_json = json.dumps(metadata, ensure_ascii=False) if metadata else "{}"
|
||||
|
||||
if asset:
|
||||
asset.status = status
|
||||
asset.local_path = local_path
|
||||
asset.remote_url = remote_url
|
||||
asset.task_id = task_id
|
||||
asset.metadata_json = meta_json
|
||||
asset.updated_at = time.time()
|
||||
else:
|
||||
new_asset = SceneAsset(
|
||||
project_id=project_id,
|
||||
scene_id=scene_id,
|
||||
asset_type=asset_type,
|
||||
status=status,
|
||||
local_path=local_path,
|
||||
remote_url=remote_url,
|
||||
task_id=task_id,
|
||||
metadata_json=meta_json,
|
||||
created_at=time.time(),
|
||||
updated_at=time.time()
|
||||
)
|
||||
session.add(new_asset)
|
||||
|
||||
session.commit()
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
logger.error(f"Error saving asset: {e}")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def get_assets(self, project_id: str, asset_type: str = None) -> List[Dict[str, Any]]:
|
||||
session = self._get_session()
|
||||
try:
|
||||
query = session.query(SceneAsset).filter_by(project_id=project_id)
|
||||
if asset_type:
|
||||
query = query.filter_by(asset_type=asset_type)
|
||||
|
||||
assets = query.all()
|
||||
results = []
|
||||
for a in assets:
|
||||
data = {
|
||||
"id": a.id,
|
||||
"project_id": a.project_id,
|
||||
"scene_id": a.scene_id,
|
||||
"asset_type": a.asset_type,
|
||||
"status": a.status,
|
||||
"local_path": a.local_path,
|
||||
"remote_url": a.remote_url,
|
||||
"task_id": a.task_id,
|
||||
"metadata": json.loads(a.metadata_json) if a.metadata_json else {},
|
||||
"updated_at": a.updated_at
|
||||
}
|
||||
results.append(data)
|
||||
return results
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def get_asset(self, project_id: str, scene_id: int, asset_type: str) -> Optional[Dict[str, Any]]:
|
||||
session = self._get_session()
|
||||
try:
|
||||
a = session.query(SceneAsset).filter_by(
|
||||
project_id=project_id,
|
||||
scene_id=scene_id,
|
||||
asset_type=asset_type
|
||||
).first()
|
||||
|
||||
if a:
|
||||
return {
|
||||
"id": a.id,
|
||||
"project_id": a.project_id,
|
||||
"scene_id": a.scene_id,
|
||||
"asset_type": a.asset_type,
|
||||
"status": a.status,
|
||||
"local_path": a.local_path,
|
||||
"remote_url": a.remote_url,
|
||||
"task_id": a.task_id,
|
||||
"metadata": json.loads(a.metadata_json) if a.metadata_json else {},
|
||||
"updated_at": a.updated_at
|
||||
}
|
||||
return None
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
# --- Config/Prompt Operations ---
|
||||
|
||||
def get_config(self, key: str, default: Any = None) -> Any:
|
||||
session = self._get_session()
|
||||
try:
|
||||
cfg = session.query(AppConfig).filter_by(key=key).first()
|
||||
if cfg:
|
||||
try:
|
||||
return json.loads(cfg.value)
|
||||
except:
|
||||
return cfg.value
|
||||
return default
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def set_config(self, key: str, value: Any, description: str = None):
|
||||
session = self._get_session()
|
||||
try:
|
||||
json_val = json.dumps(value, ensure_ascii=False)
|
||||
|
||||
cfg = session.query(AppConfig).filter_by(key=key).first()
|
||||
if cfg:
|
||||
cfg.value = json_val
|
||||
if description:
|
||||
cfg.description = description
|
||||
cfg.updated_at = time.time()
|
||||
else:
|
||||
new_cfg = AppConfig(
|
||||
key=key,
|
||||
value=json_val,
|
||||
description=description,
|
||||
updated_at=time.time()
|
||||
)
|
||||
session.add(new_cfg)
|
||||
session.commit()
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
logger.error(f"Error setting config: {e}")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
# Singleton instance
|
||||
db = DBManager()
|
||||
Reference in New Issue
Block a user