Sepehr 550f3516db feat: Add PostgreSQL database infrastructure
- 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
2025-12-31 10:56:19 +01:00

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'),
)