kunthawat/ALwrity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Recovered state: integrated TrendSurferAgent, restored frontend/backend files, and cleaned up recovery scripts

Browse Source
This commit is contained in:
ajaysi
2026-02-08 13:56:57 +05:30
parent 1db10ccd0f
commit e404a86502
333 changed files with 42223 additions and 10875 deletions
Download Patch File Download Diff File Expand all files Collapse all files

252
backend/services/database.py
View File

@@ -8,7 +8,7 @@ from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.exc import SQLAlchemyError
from loguru import logger
from typing import Optional
from typing import Optional, List
# Import models
from models.onboarding import Base as OnboardingBase
@@ -17,6 +17,7 @@ from models.content_planning import Base as ContentPlanningBase
from models.enhanced_strategy_models import Base as EnhancedStrategyBase
# Monitoring models now use the same base as enhanced strategy models
from models.monitoring_models import Base as MonitoringBase
from models.api_monitoring import Base as APIMonitoringBase
from models.persona_models import Base as PersonaBase
from models.subscription_models import Base as SubscriptionBase
from models.user_business_info import Base as UserBusinessInfoBase
@@ -27,50 +28,94 @@ from models.product_marketing_models import Campaign, CampaignProposal, Campaign
from models.product_asset_models import ProductAsset, ProductStyleTemplate, EcommerceExport
# Podcast Maker models use SubscriptionBase, but import to ensure models are registered
from models.podcast_models import PodcastProject
# Research models use SubscriptionBase
from models.research_models import ResearchProject
# Bing Analytics models
from models.bing_analytics_models import Base as BingAnalyticsBase
# Monitoring Task Models (Share EnhancedStrategyBase but need explicit import to register)
# Import these to ensure their tables are created by EnhancedStrategyBase.metadata.create_all
import models.oauth_token_monitoring_models
import models.website_analysis_monitoring_models
import models.platform_insights_monitoring_models
import models.agent_activity_models
import models.daily_workflow_models
# Database configuration
DATABASE_URL = os.getenv('DATABASE_URL', 'sqlite:///./alwrity.db')
# Get project root (3 levels up from services/database.py: services -> backend -> root)
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
WORKSPACE_DIR = os.path.join(ROOT_DIR, 'workspace')
# Create engine with safer pooling defaults and SQLite-friendly settings
engine_kwargs = {
"echo": False, # Set to True for SQL debugging
"pool_pre_ping": True, # Detect stale connections
"pool_recycle": 300, # Recycle connections to avoid timeouts
"pool_size": int(os.getenv("DB_POOL_SIZE", "20")),
"max_overflow": int(os.getenv("DB_MAX_OVERFLOW", "40")),
"pool_timeout": int(os.getenv("DB_POOL_TIMEOUT", "30")),
}
# Engine cache for multi-tenant support
_user_engines = {}
# SQLite needs special handling for multithreaded FastAPI
if DATABASE_URL.startswith("sqlite"):
engine_kwargs["connect_args"] = {"check_same_thread": False}
def get_user_db_path(user_id: str) -> str:
"""Get the database path for a specific user."""
# Sanitize user_id to be safe for filesystem
safe_user_id = "".join(c for c in user_id if c.isalnum() or c in ('-', '_'))
user_workspace = os.path.join(WORKSPACE_DIR, f"workspace_{safe_user_id}")
return os.path.join(user_workspace, 'db', f'alwrity_{safe_user_id}.db')
engine = create_engine(
DATABASE_URL,
**engine_kwargs,
)
# Create session factory
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def get_db_session() -> Optional[Session]:
def get_all_user_ids() -> List[str]:
"""
Get a database session.
Discover all user IDs by scanning workspace directories.
Returns a list of user_ids (e.g., 'user_2p...', 'user_123').
"""
user_ids = []
if not os.path.exists(WORKSPACE_DIR):
return []
Returns:
Database session or None if connection fails
"""
try:
db = SessionLocal()
return db
except SQLAlchemyError as e:
logger.error(f"Error creating database session: {str(e)}")
return None
for item in os.listdir(WORKSPACE_DIR):
if item.startswith("workspace_") and os.path.isdir(os.path.join(WORKSPACE_DIR, item)):
# Extract user_id from workspace_{user_id}
user_id = item[len("workspace_"):]
if user_id:
user_ids.append(user_id)
except Exception as e:
logger.error(f"Error discovering user workspaces: {e}")
return user_ids
def init_database():
"""
Initialize the database by creating all tables.
"""
def get_engine_for_user(user_id: str):
"""Get or create a SQLAlchemy engine for a specific user."""
if user_id in _user_engines:
return _user_engines[user_id]
db_path = get_user_db_path(user_id)
os.makedirs(os.path.dirname(db_path), exist_ok=True)
database_url = f"sqlite:///{db_path}"
engine_kwargs = {
"echo": False,
"pool_pre_ping": True,
"pool_recycle": 300,
"pool_size": int(os.getenv("DB_POOL_SIZE", "20")),
"max_overflow": int(os.getenv("DB_MAX_OVERFLOW", "40")),
"pool_timeout": int(os.getenv("DB_POOL_TIMEOUT", "30")),
"connect_args": {"check_same_thread": False}
}
engine = create_engine(database_url, **engine_kwargs)
_user_engines[user_id] = engine
# Ensure tables are initialized for this user
# This runs once per process per user when the engine is created
try:
# We need to import the function here or rely on it being available in the module scope
# Since this function is called at runtime, init_user_database should be available
init_user_database(user_id)
except Exception as e:
logger.error(f"Failed to auto-initialize database for user {user_id}: {e}")
# We don't raise here to allow the engine to be returned,
# but the application might fail later if tables are missing.
return engine
def init_user_database(user_id: str):
"""Initialize database tables for a specific user."""
engine = get_engine_for_user(user_id)
try:
# Create all tables for all models
OnboardingBase.metadata.create_all(bind=engine)
@@ -78,32 +123,137 @@ def init_database():
ContentPlanningBase.metadata.create_all(bind=engine)
EnhancedStrategyBase.metadata.create_all(bind=engine)
MonitoringBase.metadata.create_all(bind=engine)
APIMonitoringBase.metadata.create_all(bind=engine)
PersonaBase.metadata.create_all(bind=engine)
SubscriptionBase.metadata.create_all(bind=engine) # Includes product_marketing models
SubscriptionBase.metadata.create_all(bind=engine)
UserBusinessInfoBase.metadata.create_all(bind=engine)
ContentAssetBase.metadata.create_all(bind=engine)
logger.info("Database initialized successfully with all models including subscription system, product marketing, business info, and content assets")
# Initialize default data for new databases
try:
# Import here to avoid circular dependencies
from services.subscription.pricing_service import PricingService
# Create a session for data initialization
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
db = SessionLocal()
try:
pricing_service = PricingService(db)
pricing_service.initialize_default_pricing()
pricing_service.initialize_default_plans()
db.commit()
logger.info(f"Default pricing and plans initialized for user {user_id}")
except Exception as data_error:
logger.error(f"Error initializing default data for user {user_id}: {data_error}")
db.rollback()
finally:
db.close()
except Exception as import_error:
logger.warning(f"Could not initialize pricing data (PricingService import failed): {import_error}")
logger.info(f"Database initialized successfully for user {user_id}")
except SQLAlchemyError as e:
logger.error(f"Error initializing database: {str(e)}")
logger.error(f"Error initializing database for user {user_id}: {str(e)}")
raise
def init_database():
"""
Initialize global database tables (for backward compatibility/startup checks).
Uses default engine.
"""
if not default_engine:
logger.warning("Global database initialization skipped: default_engine is disabled (Multi-tenant mode)")
return
try:
# Create all tables for all models using default engine
OnboardingBase.metadata.create_all(bind=default_engine)
SEOAnalysisBase.metadata.create_all(bind=default_engine)
ContentPlanningBase.metadata.create_all(bind=default_engine)
EnhancedStrategyBase.metadata.create_all(bind=default_engine)
MonitoringBase.metadata.create_all(bind=default_engine)
APIMonitoringBase.metadata.create_all(bind=default_engine)
PersonaBase.metadata.create_all(bind=default_engine)
SubscriptionBase.metadata.create_all(bind=default_engine)
UserBusinessInfoBase.metadata.create_all(bind=default_engine)
ContentAssetBase.metadata.create_all(bind=default_engine)
logger.info("Global database initialized successfully")
except SQLAlchemyError as e:
logger.error(f"Error initializing global database: {str(e)}")
# Import here to avoid circular dependency at module level if possible,
# but get_db needs it.
# We assume auth_middleware is available.
from middleware.auth_middleware import get_current_user
from fastapi import Depends
# Legacy support for single-tenant code
# TODO: Refactor all consumers to use get_db or get_session_for_user
default_db_path = None # os.path.join(ROOT_DIR, 'alwrity.db')
DATABASE_URL = None # f"sqlite:///{default_db_path}"
default_engine = None # create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
engine = None # default_engine
SessionLocal = None # sessionmaker(autocommit=False, autoflush=False, bind=default_engine)
def get_db(current_user: dict = Depends(get_current_user)):
"""
Database dependency for FastAPI endpoints.
Context-aware: connects to the authenticated user's database.
"""
user_id = current_user.get('id') or current_user.get('clerk_user_id')
if not user_id:
# Fallback or error? For now log error
logger.error("No user ID found in context for DB connection")
# Could raise exception, but let's try to be safe
raise Exception("User ID required for database access")
engine = get_engine_for_user(user_id)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
db = SessionLocal()
try:
yield db
finally:
db.close()
# Helper for scripts/legacy that explicitly know the user_id
def get_session_for_user(user_id: str) -> Optional[Session]:
"""
Get a new database session for a specific user.
The session is not scoped, so the caller is responsible for closing it.
"""
engine = get_engine_for_user(user_id)
if not engine:
return None
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
return SessionLocal()
def get_db_session(user_id: Optional[str] = None) -> Optional[Session]:
"""
DEPRECATED: Use get_session_for_user(user_id) instead.
Legacy wrapper to prevent ImportErrors during refactoring.
"""
from utils.logger_utils import get_service_logger
logger = get_service_logger("database")
# logger.warning("Using deprecated get_db_session. Please update to get_session_for_user(user_id).")
if user_id:
return get_session_for_user(user_id)
# If no user_id, we can't give a valid session in multi-tenant mode
return None
def close_database():
"""
Close database connections.
"""
try:
engine.dispose()
for engine in _user_engines.values():
engine.dispose()
_user_engines.clear()
logger.info("Database connections closed")
except Exception as e:
logger.error(f"Error closing database connections: {str(e)}")
# Database dependency for FastAPI
def get_db():
"""
Database dependency for FastAPI endpoints.
"""
db = SessionLocal()
try:
yield db
finally:
db.close()