Fix types in database package

src/config.py

@@ -20,10 +20,10 @@ from typing import TYPE_CHECKING, ClassVar, Optional
 # Third-Party Library Imports
 from dotenv import load_dotenv
 
+# Local Application Imports
 if TYPE_CHECKING:
     from src.integrations import AnthropicConfig, ElevenLabsConfig, HumeConfig
 
-
 logger: logging.Logger = logging.getLogger("tts_arena")
 
 

src/database/crud.py

@@ -32,7 +32,7 @@ def create_vote(db: Session, vote_data: VotingResults) -> VoteResult:
         option_b_provider=vote_data["option_b_provider"],
         option_a_generation_id=vote_data["option_a_generation_id"],
         option_b_generation_id=vote_data["option_b_generation_id"],
-        voice_description=vote_data["voice_description"],
+        voice_description=vote_data["character_description"],
         text=vote_data["text"],
         is_custom_text=vote_data["is_custom_text"],
     )

src/database/database.py

@@ -5,7 +5,7 @@ This module sets up the SQLAlchemy database connection for the Expressive TTS Ar
 It initializes the PostgreSQL engine, creates a session factory for handling database transactions,
 and defines a declarative base class for ORM models.
 
-If no DATABASE_URL environment variable is set, then create a dummy database to fail gracefully
+If no DATABASE_URL environment variable is set, then create a dummy database to fail gracefully.
 """
 
 # Standard Library Imports
@@ -13,12 +13,20 @@ from typing import Callable, Optional
 
 # Third-Party Library Imports
 from sqlalchemy import Engine, create_engine
-from sqlalchemy.orm import declarative_base, sessionmaker
+from sqlalchemy.orm import DeclarativeBase, sessionmaker
 
 # Local Application Imports
 from src.config import Config
 
 
+# Define the SQLAlchemy Base using SQLAlchemy 2.0 style.
+class Base(DeclarativeBase):
+    pass
+
+
+engine: Optional[Engine] = None
+
+
 class DummySession:
     is_dummy = True  # Flag to indicate this is a dummy session.
 
@@ -49,32 +57,42 @@ class DummySession:
         pass
 
 
-Base = declarative_base()
-engine: Optional[Engine] = None
-
-
+# DBSessionMaker is either a sessionmaker instance or a callable that returns a DummySession.
 DBSessionMaker = sessionmaker | Callable[[], DummySession]
 
 
 def init_db(config: Config) -> DBSessionMaker:
+    """
+    Initialize the database engine and return a session factory based on the provided configuration.
+
+    In production, a valid DATABASE_URL is required. In development, if a DATABASE_URL is provided,
+    it is used; otherwise, a dummy session factory is returned to allow graceful failure.
+
+    Args:
+        config (Config): The application configuration.
+
+    Returns:
+        DBSessionMaker: A sessionmaker bound to the engine, or a dummy session factory.
+    """
     # ruff doesn't like setting global variables, but this is practical here
-    global engine # noqa 
+    global engine  # noqa
 
     if config.app_env == "prod":
         # In production, a valid DATABASE_URL is required.
         if not config.database_url:
             raise ValueError("DATABASE_URL must be set in production!")
-
         engine = create_engine(config.database_url)
         return sessionmaker(bind=engine)
+
     # In development, if a DATABASE_URL is provided, use it.
     if config.database_url:
         engine = create_engine(config.database_url)
         return sessionmaker(bind=engine)
+
     # No DATABASE_URL is provided; use a DummySession that does nothing.
     engine = None
 
-    def dummy_session_factory():
+    def dummy_session_factory() -> DummySession:
         return DummySession()
 
     return dummy_session_factory

src/database/models.py

@@ -42,7 +42,7 @@ class VoteResult(Base):
     created_at = Column(DateTime, nullable=False, server_default=func.now())
     comparison_type = Column(String(50), nullable=False)
     winning_provider = Column(String(50), nullable=False)
-    winning_option = Column(saEnum(OptionEnum), nullable=False)
+    winning_option = Column(saEnum(OptionEnum), nullable=False)  # type: ignore
     option_a_provider = Column(String(50), nullable=False)
     option_b_provider = Column(String(50), nullable=False)
     option_a_generation_id = Column(String(100), nullable=True)