149 lines
5.4 KiB
Python
149 lines
5.4 KiB
Python
"""Contains all SQL code here."""
|
|
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
import env
|
|
from aiologger import Logger
|
|
from models import Genre
|
|
from models import Statistics
|
|
from models import Track
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlalchemy.future import select
|
|
from sqlalchemy.orm import registry
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
|
|
# Initialize asynchronous logger
|
|
logger = Logger.with_default_handlers(name="sql_logger")
|
|
|
|
engine = create_async_engine(env.DATABASE_URL, echo=True)
|
|
|
|
|
|
mapper_registry = registry()
|
|
|
|
async_session_maker = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
async def init_db() -> None:
|
|
"""Create all tables asynchronously."""
|
|
try:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(mapper_registry.metadata.create_all)
|
|
await logger.info("Database tables created.")
|
|
except Exception as e:
|
|
await logger.error(f"Failed to create tables: {e}")
|
|
|
|
|
|
async def drop_db() -> None:
|
|
"""Drop all tables asynchronously for clean slate testing or teardown."""
|
|
try:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(mapper_registry.metadata.drop_all)
|
|
await logger.info("Database tables dropped.")
|
|
except Exception as e:
|
|
await logger.error(f"Failed to drop tables: {e}")
|
|
|
|
|
|
async def insert_or_update_track(track_data: dict[str, Any]) -> None:
|
|
"""Insert a new track or update an existing one in the database asynchronously."""
|
|
async with async_session_maker() as session:
|
|
try:
|
|
async with session.begin():
|
|
genre = await session.get(Genre, track_data.get("genre_id"))
|
|
if not genre:
|
|
genre = Genre(name=track_data["genre_name"])
|
|
session.add(genre)
|
|
await session.flush() # Ensures 'genre' is persisted and has an 'id'
|
|
|
|
track = await session.get(Track, track_data.get("id"))
|
|
if track:
|
|
for key, value in track_data.items():
|
|
setattr(track, key, value)
|
|
await logger.info(f"Updated track: {track.title}")
|
|
else:
|
|
track = Track(**track_data)
|
|
session.add(track)
|
|
await logger.info(f"Inserted new track: {track.title}")
|
|
|
|
await session.commit()
|
|
except SQLAlchemyError as e:
|
|
await logger.error(f"Database error in insert_or_update_track: {e}")
|
|
await session.rollback()
|
|
except Exception as e:
|
|
await logger.error(f"Unexpected error in insert_or_update_track: {e}")
|
|
await session.rollback()
|
|
|
|
|
|
async def cleanup_unused_genres() -> None:
|
|
"""Remove genres that are no longer used by any tracks."""
|
|
async with async_session_maker() as session:
|
|
try:
|
|
async with session.begin():
|
|
stmt = select(Genre).outerjoin(Track).filter(Track.id is None)
|
|
result = await session.execute(stmt)
|
|
unused_genres = result.scalars().all()
|
|
|
|
for genre in unused_genres:
|
|
await session.delete(genre)
|
|
|
|
await session.commit()
|
|
await logger.info(f"Cleaned up {len(unused_genres)} unused genres.")
|
|
except SQLAlchemyError as e:
|
|
await logger.error(f"Database error in cleanup_unused_genres: {e}")
|
|
await session.rollback()
|
|
except Exception as e:
|
|
await logger.error(f"Unexpected error in cleanup_unused_genres: {e}")
|
|
await session.rollback()
|
|
|
|
|
|
async def update_statistics_timestamps(operation: str) -> None:
|
|
"""Update the start or end timestamps in the statistics model based on the specified operation.
|
|
|
|
Args:
|
|
operation (str): A string to specify whether to update 'start' or 'end' timestamp.
|
|
"""
|
|
async with async_session_maker() as session:
|
|
try:
|
|
async with session.begin():
|
|
stats = await session.execute(
|
|
select(Statistics).order_by(Statistics.id)
|
|
)
|
|
stats_obj = stats.scalars().first()
|
|
|
|
if operation == "start":
|
|
stats_obj.track_update_start_timestamp = datetime.now()
|
|
stats_obj.track_update_end_timestamp = None
|
|
log_message = "Statistics start timestamp updated."
|
|
elif operation == "end":
|
|
stats_obj.track_update_end_timestamp = datetime.now()
|
|
log_message = "Statistics end timestamp updated."
|
|
else:
|
|
raise ValueError(
|
|
"Invalid operation specified for statistics timestamp update."
|
|
)
|
|
|
|
await session.commit()
|
|
await logger.info(log_message)
|
|
|
|
except SQLAlchemyError as e:
|
|
await logger.error(
|
|
f"Database error during statistics timestamp update: {e}"
|
|
)
|
|
await session.rollback()
|
|
except Exception as e:
|
|
await logger.error(
|
|
f"Unexpected error during statistics timestamp update: {e}"
|
|
)
|
|
await session.rollback()
|
|
|
|
|
|
# Example usage
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
|
|
asyncio.run(update_statistics_timestamps("start"))
|
|
asyncio.run(update_statistics_timestamps("end"))
|