- Add SQLAlchemy models for User, Translation, ApiKey, UsageLog, PaymentHistory - Add database connection management with PostgreSQL/SQLite support - Add repository layer for CRUD operations - Add Alembic migration setup with initial migration - Update auth_service to automatically use database when DATABASE_URL is set - Update docker-compose.yml with PostgreSQL service and Redis (non-optional) - Add database migration script (scripts/migrate_to_db.py) - Update .env.example with database configuration
260 lines
9.0 KiB
Python
260 lines
9.0 KiB
Python
"""
|
|
SQLAlchemy models for the Document Translation API
|
|
"""
|
|
import os
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Optional, List
|
|
|
|
from sqlalchemy import (
|
|
Column, String, Integer, Float, Boolean, DateTime, Text,
|
|
ForeignKey, Enum, Index, JSON, BigInteger
|
|
)
|
|
from sqlalchemy.orm import relationship, declarative_base
|
|
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
|
|
import enum
|
|
|
|
Base = declarative_base()
|
|
|
|
|
|
def generate_uuid():
|
|
"""Generate a new UUID string"""
|
|
return str(uuid.uuid4())
|
|
|
|
|
|
class PlanType(str, enum.Enum):
|
|
FREE = "free"
|
|
STARTER = "starter"
|
|
PRO = "pro"
|
|
BUSINESS = "business"
|
|
ENTERPRISE = "enterprise"
|
|
|
|
|
|
class SubscriptionStatus(str, enum.Enum):
|
|
ACTIVE = "active"
|
|
CANCELED = "canceled"
|
|
PAST_DUE = "past_due"
|
|
TRIALING = "trialing"
|
|
PAUSED = "paused"
|
|
|
|
|
|
class User(Base):
|
|
"""User model for authentication and billing"""
|
|
__tablename__ = "users"
|
|
|
|
id = Column(String(36), primary_key=True, default=generate_uuid)
|
|
email = Column(String(255), unique=True, nullable=False, index=True)
|
|
name = Column(String(255), nullable=False)
|
|
password_hash = Column(String(255), nullable=False)
|
|
|
|
# Account status
|
|
email_verified = Column(Boolean, default=False)
|
|
is_active = Column(Boolean, default=True)
|
|
avatar_url = Column(String(500), nullable=True)
|
|
|
|
# Subscription info
|
|
plan = Column(Enum(PlanType), default=PlanType.FREE)
|
|
subscription_status = Column(Enum(SubscriptionStatus), default=SubscriptionStatus.ACTIVE)
|
|
|
|
# Stripe integration
|
|
stripe_customer_id = Column(String(255), nullable=True, index=True)
|
|
stripe_subscription_id = Column(String(255), nullable=True)
|
|
|
|
# Usage tracking (reset monthly)
|
|
docs_translated_this_month = Column(Integer, default=0)
|
|
pages_translated_this_month = Column(Integer, default=0)
|
|
api_calls_this_month = Column(Integer, default=0)
|
|
extra_credits = Column(Integer, default=0) # Purchased credits
|
|
usage_reset_date = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Timestamps
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
|
last_login_at = Column(DateTime, nullable=True)
|
|
|
|
# Relationships
|
|
translations = relationship("Translation", back_populates="user", lazy="dynamic")
|
|
api_keys = relationship("ApiKey", back_populates="user", lazy="dynamic")
|
|
|
|
# Indexes
|
|
__table_args__ = (
|
|
Index('ix_users_email_active', 'email', 'is_active'),
|
|
Index('ix_users_stripe_customer', 'stripe_customer_id'),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Convert user to dictionary for API response"""
|
|
return {
|
|
"id": self.id,
|
|
"email": self.email,
|
|
"name": self.name,
|
|
"avatar_url": self.avatar_url,
|
|
"plan": self.plan.value if self.plan else "free",
|
|
"subscription_status": self.subscription_status.value if self.subscription_status else "active",
|
|
"docs_translated_this_month": self.docs_translated_this_month,
|
|
"pages_translated_this_month": self.pages_translated_this_month,
|
|
"api_calls_this_month": self.api_calls_this_month,
|
|
"extra_credits": self.extra_credits,
|
|
"email_verified": self.email_verified,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
}
|
|
|
|
|
|
class Translation(Base):
|
|
"""Translation history for analytics and billing"""
|
|
__tablename__ = "translations"
|
|
|
|
id = Column(String(36), primary_key=True, default=generate_uuid)
|
|
user_id = Column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
|
|
|
|
# File info
|
|
original_filename = Column(String(255), nullable=False)
|
|
file_type = Column(String(10), nullable=False) # xlsx, docx, pptx
|
|
file_size_bytes = Column(BigInteger, default=0)
|
|
page_count = Column(Integer, default=0)
|
|
|
|
# Translation details
|
|
source_language = Column(String(10), default="auto")
|
|
target_language = Column(String(10), nullable=False)
|
|
provider = Column(String(50), nullable=False) # google, deepl, ollama, etc.
|
|
|
|
# Processing info
|
|
status = Column(String(20), default="pending") # pending, processing, completed, failed
|
|
error_message = Column(Text, nullable=True)
|
|
processing_time_ms = Column(Integer, nullable=True)
|
|
|
|
# Cost tracking (for paid providers)
|
|
characters_translated = Column(Integer, default=0)
|
|
estimated_cost_usd = Column(Float, default=0.0)
|
|
|
|
# Timestamps
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
completed_at = Column(DateTime, nullable=True)
|
|
|
|
# Relationship
|
|
user = relationship("User", back_populates="translations")
|
|
|
|
# Indexes
|
|
__table_args__ = (
|
|
Index('ix_translations_user_date', 'user_id', 'created_at'),
|
|
Index('ix_translations_status', 'status'),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"id": self.id,
|
|
"original_filename": self.original_filename,
|
|
"file_type": self.file_type,
|
|
"file_size_bytes": self.file_size_bytes,
|
|
"page_count": self.page_count,
|
|
"source_language": self.source_language,
|
|
"target_language": self.target_language,
|
|
"provider": self.provider,
|
|
"status": self.status,
|
|
"processing_time_ms": self.processing_time_ms,
|
|
"characters_translated": self.characters_translated,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
|
|
}
|
|
|
|
|
|
class ApiKey(Base):
|
|
"""API keys for programmatic access"""
|
|
__tablename__ = "api_keys"
|
|
|
|
id = Column(String(36), primary_key=True, default=generate_uuid)
|
|
user_id = Column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
|
|
|
|
# Key info
|
|
name = Column(String(100), nullable=False) # User-friendly name
|
|
key_hash = Column(String(255), nullable=False) # SHA256 of the key
|
|
key_prefix = Column(String(10), nullable=False) # First 8 chars for identification
|
|
|
|
# Permissions
|
|
is_active = Column(Boolean, default=True)
|
|
scopes = Column(JSON, default=list) # ["translate", "read", "write"]
|
|
|
|
# Usage tracking
|
|
last_used_at = Column(DateTime, nullable=True)
|
|
usage_count = Column(Integer, default=0)
|
|
|
|
# Timestamps
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
expires_at = Column(DateTime, nullable=True)
|
|
|
|
# Relationship
|
|
user = relationship("User", back_populates="api_keys")
|
|
|
|
# Indexes
|
|
__table_args__ = (
|
|
Index('ix_api_keys_prefix', 'key_prefix'),
|
|
Index('ix_api_keys_hash', 'key_hash'),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"id": self.id,
|
|
"name": self.name,
|
|
"key_prefix": self.key_prefix,
|
|
"is_active": self.is_active,
|
|
"scopes": self.scopes,
|
|
"last_used_at": self.last_used_at.isoformat() if self.last_used_at else None,
|
|
"usage_count": self.usage_count,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
"expires_at": self.expires_at.isoformat() if self.expires_at else None,
|
|
}
|
|
|
|
|
|
class UsageLog(Base):
|
|
"""Daily usage aggregation for billing and analytics"""
|
|
__tablename__ = "usage_logs"
|
|
|
|
id = Column(String(36), primary_key=True, default=generate_uuid)
|
|
user_id = Column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
|
|
|
|
# Date (for daily aggregation)
|
|
date = Column(DateTime, nullable=False, index=True)
|
|
|
|
# Aggregated counts
|
|
documents_count = Column(Integer, default=0)
|
|
pages_count = Column(Integer, default=0)
|
|
characters_count = Column(BigInteger, default=0)
|
|
api_calls_count = Column(Integer, default=0)
|
|
|
|
# By provider breakdown (JSON)
|
|
provider_breakdown = Column(JSON, default=dict)
|
|
|
|
# Indexes
|
|
__table_args__ = (
|
|
Index('ix_usage_logs_user_date', 'user_id', 'date', unique=True),
|
|
)
|
|
|
|
|
|
class PaymentHistory(Base):
|
|
"""Payment and invoice history"""
|
|
__tablename__ = "payment_history"
|
|
|
|
id = Column(String(36), primary_key=True, default=generate_uuid)
|
|
user_id = Column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
|
|
|
|
# Stripe info
|
|
stripe_payment_intent_id = Column(String(255), nullable=True)
|
|
stripe_invoice_id = Column(String(255), nullable=True)
|
|
|
|
# Payment details
|
|
amount_cents = Column(Integer, nullable=False)
|
|
currency = Column(String(3), default="usd")
|
|
payment_type = Column(String(50), nullable=False) # subscription, credits, one_time
|
|
status = Column(String(20), nullable=False) # succeeded, failed, pending, refunded
|
|
|
|
# Description
|
|
description = Column(String(255), nullable=True)
|
|
|
|
# Timestamps
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Indexes
|
|
__table_args__ = (
|
|
Index('ix_payment_history_user', 'user_id', 'created_at'),
|
|
)
|