A kind of initial commit
This commit is contained in:
heboba
0
app/db/__init__.py
Normal file
0
app/db/__init__.py
Normal file
18
app/db/base.py
Normal file
18
app/db/base.py
Normal file
@@ -0,0 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
|
||||
def utcnow() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
class TimestampMixin:
|
||||
created_at: Mapped[datetime] = mapped_column(default=utcnow)
|
||||
updated_at: Mapped[datetime] = mapped_column(default=utcnow, onupdate=utcnow)
|
||||
115
app/db/models.py
Normal file
115
app/db/models.py
Normal file
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
|
||||
from sqlalchemy import Boolean, Date, DateTime, Float, ForeignKey, Integer, String, Text, UniqueConstraint
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from app.db.base import Base, TimestampMixin
|
||||
|
||||
|
||||
class User(Base, TimestampMixin):
|
||||
__tablename__ = "users"
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
telegram_chat_id: Mapped[str] = mapped_column(String(64), unique=True, index=True)
|
||||
telegram_username: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
|
||||
spotify_user_id: Mapped[str | None] = mapped_column(String(128), unique=True, nullable=True, index=True)
|
||||
spotify_access_token: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
spotify_refresh_token: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
spotify_token_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
spotify_scopes: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
|
||||
timezone: Mapped[str] = mapped_column(String(64), default="UTC")
|
||||
is_active: Mapped[bool] = mapped_column(Boolean, default=True)
|
||||
playlist_size: Mapped[int] = mapped_column(Integer, default=30)
|
||||
min_new_ratio: Mapped[float] = mapped_column(Float, default=0.8)
|
||||
last_generated_date: Mapped[date | None] = mapped_column(Date, nullable=True)
|
||||
latest_playlist_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
latest_playlist_url: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
|
||||
saved_tracks: Mapped[list["SavedTrack"]] = relationship(back_populates="user", cascade="all, delete-orphan")
|
||||
playlist_runs: Mapped[list["PlaylistRun"]] = relationship(back_populates="user", cascade="all, delete-orphan")
|
||||
rec_history: Mapped[list["RecommendationHistory"]] = relationship(
|
||||
back_populates="user", cascade="all, delete-orphan"
|
||||
)
|
||||
|
||||
|
||||
class AuthState(Base):
|
||||
__tablename__ = "auth_states"
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
state: Mapped[str] = mapped_column(String(128), unique=True, index=True)
|
||||
telegram_chat_id: Mapped[str] = mapped_column(String(64), index=True)
|
||||
expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
|
||||
|
||||
class SavedTrack(Base):
|
||||
__tablename__ = "saved_tracks"
|
||||
__table_args__ = (UniqueConstraint("user_id", "spotify_track_id", name="uq_saved_track_user_track"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"), index=True)
|
||||
spotify_track_id: Mapped[str] = mapped_column(String(128), index=True)
|
||||
name: Mapped[str] = mapped_column(Text)
|
||||
artist_names: Mapped[str] = mapped_column(Text)
|
||||
artist_ids_csv: Mapped[str] = mapped_column(Text)
|
||||
album_name: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
added_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
popularity: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
||||
|
||||
user: Mapped[User] = relationship(back_populates="saved_tracks")
|
||||
|
||||
|
||||
class RecommendationHistory(Base):
|
||||
__tablename__ = "recommendation_history"
|
||||
__table_args__ = (UniqueConstraint("user_id", "spotify_track_id", name="uq_rec_history_user_track"),)
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"), index=True)
|
||||
spotify_track_id: Mapped[str] = mapped_column(String(128), index=True)
|
||||
first_recommended_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
last_recommended_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
times_recommended: Mapped[int] = mapped_column(Integer, default=1)
|
||||
|
||||
user: Mapped[User] = relationship(back_populates="rec_history")
|
||||
|
||||
|
||||
class PlaylistRun(Base, TimestampMixin):
|
||||
__tablename__ = "playlist_runs"
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"), index=True)
|
||||
run_date: Mapped[date] = mapped_column(Date)
|
||||
status: Mapped[str] = mapped_column(String(32), default="pending")
|
||||
playlist_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
playlist_name: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
playlist_url: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
total_tracks: Mapped[int] = mapped_column(Integer, default=0)
|
||||
new_tracks: Mapped[int] = mapped_column(Integer, default=0)
|
||||
reused_tracks: Mapped[int] = mapped_column(Integer, default=0)
|
||||
notes: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
|
||||
user: Mapped[User] = relationship(back_populates="playlist_runs")
|
||||
tracks: Mapped[list["PlaylistRunTrack"]] = relationship(back_populates="run", cascade="all, delete-orphan")
|
||||
|
||||
|
||||
class PlaylistRunTrack(Base):
|
||||
__tablename__ = "playlist_run_tracks"
|
||||
__table_args__ = (
|
||||
UniqueConstraint("run_id", "spotify_track_id", name="uq_run_track"),
|
||||
UniqueConstraint("run_id", "position", name="uq_run_position"),
|
||||
)
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
run_id: Mapped[int] = mapped_column(ForeignKey("playlist_runs.id", ondelete="CASCADE"), index=True)
|
||||
spotify_track_id: Mapped[str] = mapped_column(String(128), index=True)
|
||||
name: Mapped[str] = mapped_column(Text)
|
||||
artist_names: Mapped[str] = mapped_column(Text)
|
||||
source: Mapped[str] = mapped_column(String(64))
|
||||
position: Mapped[int] = mapped_column(Integer)
|
||||
is_new_to_bot: Mapped[bool] = mapped_column(Boolean, default=True)
|
||||
|
||||
run: Mapped[PlaylistRun] = relationship(back_populates="tracks")
|
||||
201
app/db/repositories.py
Normal file
201
app/db/repositories.py
Normal file
@@ -0,0 +1,201 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from datetime import date, datetime, timezone
|
||||
|
||||
from sqlalchemy import delete, func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.db.models import (
|
||||
AuthState,
|
||||
PlaylistRun,
|
||||
PlaylistRunTrack,
|
||||
RecommendationHistory,
|
||||
SavedTrack,
|
||||
User,
|
||||
)
|
||||
from app.utils.time import ensure_utc
|
||||
|
||||
|
||||
class UserRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get_or_create_by_chat(self, chat_id: str, username: str | None = None) -> User:
|
||||
user = await self.get_by_chat_id(chat_id)
|
||||
if user:
|
||||
if username and user.telegram_username != username:
|
||||
user.telegram_username = username
|
||||
return user
|
||||
user = User(telegram_chat_id=chat_id, telegram_username=username)
|
||||
self.session.add(user)
|
||||
await self.session.flush()
|
||||
return user
|
||||
|
||||
async def get_by_chat_id(self, chat_id: str) -> User | None:
|
||||
result = await self.session.execute(select(User).where(User.telegram_chat_id == chat_id))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def get_by_id(self, user_id: int) -> User | None:
|
||||
result = await self.session.execute(select(User).where(User.id == user_id))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list_active_connected_users(self) -> list[User]:
|
||||
result = await self.session.execute(
|
||||
select(User).where(User.is_active.is_(True), User.spotify_refresh_token.is_not(None))
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
||||
class AuthStateRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create(self, state: str, telegram_chat_id: str, expires_at: datetime) -> AuthState:
|
||||
row = AuthState(
|
||||
state=state,
|
||||
telegram_chat_id=telegram_chat_id,
|
||||
expires_at=expires_at,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return row
|
||||
|
||||
async def pop_valid(self, state: str) -> AuthState | None:
|
||||
now = datetime.now(timezone.utc)
|
||||
result = await self.session.execute(select(AuthState).where(AuthState.state == state))
|
||||
row = result.scalar_one_or_none()
|
||||
if not row:
|
||||
return None
|
||||
await self.session.delete(row)
|
||||
if ensure_utc(row.expires_at) < now:
|
||||
return None
|
||||
return row
|
||||
|
||||
async def delete_expired(self) -> int:
|
||||
result = await self.session.execute(delete(AuthState).where(AuthState.expires_at < datetime.now(timezone.utc)))
|
||||
return result.rowcount or 0
|
||||
|
||||
|
||||
class SavedTrackRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def replace_for_user(self, user_id: int, tracks: Iterable[dict]) -> None:
|
||||
await self.session.execute(delete(SavedTrack).where(SavedTrack.user_id == user_id))
|
||||
for item in tracks:
|
||||
self.session.add(
|
||||
SavedTrack(
|
||||
user_id=user_id,
|
||||
spotify_track_id=item["id"],
|
||||
name=item["name"],
|
||||
artist_names=", ".join(item["artist_names"]),
|
||||
artist_ids_csv=",".join(item["artist_ids"]),
|
||||
album_name=item.get("album_name"),
|
||||
added_at=item.get("added_at"),
|
||||
popularity=item.get("popularity"),
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
|
||||
async def list_for_user(self, user_id: int) -> list[SavedTrack]:
|
||||
result = await self.session.execute(select(SavedTrack).where(SavedTrack.user_id == user_id))
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def count_for_user(self, user_id: int) -> int:
|
||||
result = await self.session.execute(select(func.count()).select_from(SavedTrack).where(SavedTrack.user_id == user_id))
|
||||
return int(result.scalar_one())
|
||||
|
||||
|
||||
class RecommendationHistoryRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def list_track_ids(self, user_id: int) -> set[str]:
|
||||
result = await self.session.execute(
|
||||
select(RecommendationHistory.spotify_track_id).where(RecommendationHistory.user_id == user_id)
|
||||
)
|
||||
return {row[0] for row in result.all()}
|
||||
|
||||
async def mark_tracks(self, user_id: int, track_ids: list[str]) -> None:
|
||||
if not track_ids:
|
||||
return
|
||||
now = datetime.now(timezone.utc)
|
||||
result = await self.session.execute(select(RecommendationHistory).where(RecommendationHistory.user_id == user_id))
|
||||
existing = {row.spotify_track_id: row for row in result.scalars().all()}
|
||||
for track_id in track_ids:
|
||||
if track_id in existing:
|
||||
row = existing[track_id]
|
||||
row.last_recommended_at = now
|
||||
row.times_recommended += 1
|
||||
else:
|
||||
self.session.add(
|
||||
RecommendationHistory(
|
||||
user_id=user_id,
|
||||
spotify_track_id=track_id,
|
||||
first_recommended_at=now,
|
||||
last_recommended_at=now,
|
||||
times_recommended=1,
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
class PlaylistRunRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create_run(self, user_id: int, run_date: date, notes: str | None = None) -> PlaylistRun:
|
||||
run = PlaylistRun(user_id=user_id, run_date=run_date, status="running", notes=notes)
|
||||
self.session.add(run)
|
||||
await self.session.flush()
|
||||
return run
|
||||
|
||||
async def add_tracks(self, run_id: int, tracks: list[dict]) -> None:
|
||||
for idx, track in enumerate(tracks, start=1):
|
||||
self.session.add(
|
||||
PlaylistRunTrack(
|
||||
run_id=run_id,
|
||||
spotify_track_id=track["id"],
|
||||
name=track["name"],
|
||||
artist_names=", ".join(track["artist_names"]),
|
||||
source=track["source"],
|
||||
position=idx,
|
||||
is_new_to_bot=track.get("is_new_to_bot", True),
|
||||
)
|
||||
)
|
||||
await self.session.flush()
|
||||
|
||||
async def mark_success(
|
||||
self,
|
||||
run: PlaylistRun,
|
||||
*,
|
||||
playlist_id: str,
|
||||
playlist_name: str,
|
||||
playlist_url: str | None,
|
||||
total_tracks: int,
|
||||
new_tracks: int,
|
||||
reused_tracks: int,
|
||||
notes: str | None = None,
|
||||
) -> None:
|
||||
run.status = "success"
|
||||
run.playlist_id = playlist_id
|
||||
run.playlist_name = playlist_name
|
||||
run.playlist_url = playlist_url
|
||||
run.total_tracks = total_tracks
|
||||
run.new_tracks = new_tracks
|
||||
run.reused_tracks = reused_tracks
|
||||
run.notes = notes
|
||||
await self.session.flush()
|
||||
|
||||
async def mark_failed(self, run: PlaylistRun, message: str) -> None:
|
||||
run.status = "failed"
|
||||
run.notes = message
|
||||
await self.session.flush()
|
||||
|
||||
async def latest_for_user(self, user_id: int) -> PlaylistRun | None:
|
||||
result = await self.session.execute(
|
||||
select(PlaylistRun).where(PlaylistRun.user_id == user_id).order_by(PlaylistRun.created_at.desc()).limit(1)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
35
app/db/session.py
Normal file
35
app/db/session.py
Normal file
@@ -0,0 +1,35 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
||||
|
||||
from app.config import Settings
|
||||
from app.db.base import Base
|
||||
|
||||
|
||||
def create_engine(settings: Settings) -> AsyncEngine:
|
||||
return create_async_engine(settings.database_url, future=True, echo=False)
|
||||
|
||||
|
||||
def create_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
|
||||
return async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
||||
|
||||
|
||||
async def init_db(engine: AsyncEngine) -> None:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def session_scope(factory: async_sessionmaker[AsyncSession]) -> AsyncIterator[AsyncSession]:
|
||||
session = factory()
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
finally:
|
||||
await session.close()
|
||||
Reference in New Issue
Block a user