273 lines
12 KiB
Python
273 lines
12 KiB
Python
from sqlalchemy import create_engine, Column, String, Integer, Boolean, ForeignKey, select
|
||
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
|
||
from typing import Any
|
||
|
||
# 1. Define the base class for declarative models
|
||
Base = declarative_base()
|
||
|
||
# 2. Define the twitch_videos table
|
||
class TwitchVideo(Base):
|
||
__tablename__ = "twitch_videos"
|
||
|
||
id = Column(String, primary_key=True)
|
||
title = Column(String, nullable=False)
|
||
created_at = Column(String, nullable=False)
|
||
view_count = Column(Integer, nullable=False)
|
||
duration = Column(String, nullable=False)
|
||
url = Column(String, nullable=False)
|
||
thumbnail_url = Column(String, nullable=False)
|
||
game_id = Column(String, nullable=False)
|
||
game_name = Column(String, nullable=False)
|
||
stream_id = Column(String, nullable=False)
|
||
creator_name = Column(String, nullable=False)
|
||
clip_is = Column(Boolean, nullable=False, default=False)
|
||
downloaded = Column(Boolean, nullable=False, default=False)
|
||
uploaded_yt = Column(Boolean, nullable=False, default=False)
|
||
uploaded_yt_chats = Column(Boolean, nullable=False, default=False)
|
||
uploaded_yt_shorts = Column(Boolean, nullable=False, default=False)
|
||
|
||
# This now actually points to twitch_videos.id
|
||
related_video_id = Column(String, ForeignKey("twitch_videos.id"), nullable=True)
|
||
|
||
# Relationship back to YoutubeInfo (uselist=True enforces a 1-to-many structure)
|
||
youtube_info = relationship("YoutubeInfo", back_populates="twitch_video", uselist=True)
|
||
chapters = relationship("Chapter", back_populates="twitch_video", cascade="all, delete-orphan")
|
||
|
||
# 3. Define the data table
|
||
class Data(Base):
|
||
__tablename__ = "data"
|
||
|
||
# In SQLAlchemy, an Integer primary key automatically acts as AUTOINCREMENT in SQLite
|
||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
name = Column(String, nullable=False)
|
||
date_shorts = Column(String, nullable=False)
|
||
date_videos = Column(String, nullable=False)
|
||
date_videos_chats = Column(String, nullable=False)
|
||
|
||
class YoutubeInfo(Base):
|
||
__tablename__ = "youtube_info"
|
||
|
||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
twitch_video_id = Column(String, ForeignKey("twitch_videos.id"), nullable=False)
|
||
yt_id = Column(String, nullable=True)
|
||
yt_date = Column(String, nullable=True)
|
||
yt_short_id = Column(String, nullable=True)
|
||
yt_short_date = Column(String, nullable=True)
|
||
yt_chat_SS_id = Column(String, nullable=True)
|
||
yt_chat_SS_date = Column(String, nullable=True)
|
||
yt_chat_over_id = Column(String, nullable=True)
|
||
yt_chat_over_date = Column(String, nullable=True)
|
||
|
||
# Relationship linking back to the twitch video record
|
||
twitch_video = relationship("TwitchVideo", back_populates="youtube_info")
|
||
|
||
class Chapter(Base):
|
||
__tablename__ = "chapters"
|
||
|
||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
twitch_video_id = Column(String, ForeignKey("twitch_videos.id"), nullable=False)
|
||
start_milliseconds = Column(Integer, nullable=False)
|
||
length_milliseconds = Column(Integer, nullable=False)
|
||
description = Column(String, nullable=True)
|
||
type = Column(String, nullable=True)
|
||
game_display_name = Column(String, nullable=True)
|
||
game_id = Column(String, nullable=True)
|
||
|
||
# Relationship linking back to the twitch video record
|
||
twitch_video = relationship("TwitchVideo", back_populates="chapters")
|
||
|
||
class Repository:
|
||
def __init__(self, db_path: str = "database.db"):
|
||
# 4. Connect to the SQLite file and generate the tables
|
||
# This will create 'database.db' in the current working directory if it doesn't exist.
|
||
engine = create_engine(f"sqlite:///{db_path}", echo=True)
|
||
Base.metadata.create_all(engine)
|
||
|
||
# 5. Optional: Set up a session factory for interacting with the database
|
||
Session = sessionmaker(bind=engine)
|
||
self.session = Session()
|
||
|
||
def set_related_video(self, record_id: str, related_video_id: str):
|
||
record = self.session.get(TwitchVideo, record_id)
|
||
|
||
if record:
|
||
setattr(record, "related_video", related_video_id)
|
||
self.session.commit()
|
||
else:
|
||
print(f"❌ Record {record_id} not found.")
|
||
|
||
def __mark_as(self, record_id: str, field_name: str, value: bool = True):
|
||
record = self.session.get(TwitchVideo, record_id)
|
||
|
||
if record:
|
||
setattr(record, field_name, value)
|
||
self.session.commit()
|
||
else:
|
||
print(f"❌ Record {record_id} not found in __mark_as.")
|
||
|
||
def mark_as_uploaded_shorts(self, record_id: str):
|
||
"""Flags a specific row record to uploaded_yt_shorts (1)."""
|
||
self.__mark_as(record_id, "uploaded_yt_shorts")
|
||
|
||
def mark_as_uploaded_chats(self, record_id: str):
|
||
"""Flags a specific row record to uploaded_yt_chats (1)."""
|
||
self.__mark_as(record_id, "uploaded_yt_chats")
|
||
|
||
def mark_as_uploaded_yt(self, record_id: str):
|
||
"""Flags a specific row record to uploaded_yt (1)."""
|
||
self.__mark_as(record_id, "uploaded_yt")
|
||
|
||
def mark_as_downloaded(self, record_id: str):
|
||
"""Flags a specific row record to downloaded (1)."""
|
||
self.__mark_as(record_id, "downloaded")
|
||
|
||
def unmark_as_download(self, record_id: str):
|
||
"""Flags a specific row record to downloaded (0)."""
|
||
self.__mark_as(record_id, "downloaded", False)
|
||
|
||
def update_datetime(self, record_id: str, created_at: str):
|
||
record = self.session.get(TwitchVideo, record_id)
|
||
|
||
if record:
|
||
record.created_at = created_at
|
||
self.session.commit()
|
||
else:
|
||
print(f"❌ Record {record_id} not found in __mark_as.")
|
||
|
||
def __get_unuploaded(self, field_name: str, also_clause: Any = None) -> list[TwitchVideo]:
|
||
"""Retrieve all rows that were downloaded but not uploaded using SQLAlchemy."""
|
||
|
||
# 1. Validation: Verify the string corresponds to an actual column
|
||
if not hasattr(TwitchVideo, field_name):
|
||
raise ValueError(f"🚨 Invalid column name: '{field_name}' does not exist on TwitchVideo.")
|
||
|
||
# Get the actual SQLAlchemy column object dynamically
|
||
status_column = getattr(TwitchVideo, field_name)
|
||
|
||
# 2. Build the query programmatically
|
||
stmt = select(TwitchVideo).where(
|
||
TwitchVideo.downloaded == True,
|
||
status_column == False
|
||
)
|
||
|
||
# 3. Safely append additional filters if provided
|
||
if also_clause is not None:
|
||
stmt = stmt.where(also_clause)
|
||
|
||
# 4. Execute and return results
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_unuploaded_shorts(self) -> list[TwitchVideo]:
|
||
"""Retrieve all rows that were download but not uploaded_yt_shorts"""
|
||
return self.__get_unuploaded("uploaded_yt_shorts", (TwitchVideo.clip_is == True))
|
||
|
||
def get_unuploaded_chats(self) -> list[TwitchVideo]:
|
||
"""Retrieve all rows that were download but not uploaded_yt_chats"""
|
||
return self.__get_unuploaded("uploaded_yt_chats", (TwitchVideo.clip_is == False))
|
||
|
||
def get_unuploaded(self) -> list[TwitchVideo]:
|
||
"""Retrieve all rows that were download but not uploaded_yt"""
|
||
return self.__get_unuploaded("uploaded_yt")
|
||
|
||
def get_undownloaded(self) -> list[TwitchVideo]:
|
||
stmt = select(TwitchVideo).where(TwitchVideo.downloaded == False)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_downloaded(self) -> list[TwitchVideo]:
|
||
stmt = select(TwitchVideo).where(TwitchVideo.downloaded == True)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_clips(self) -> list[TwitchVideo]:
|
||
"""Retrieves all rows that are clip_is"""
|
||
stmt = select(TwitchVideo).where(TwitchVideo.clip_is == True)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_vods(self) -> list[TwitchVideo]:
|
||
stmt = select(TwitchVideo).where(TwitchVideo.clip_is == False)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def __update_date(self, field_name: str, date: str):
|
||
record_id = 1
|
||
record = self.session.get(Data, record_id)
|
||
|
||
if record:
|
||
setattr(record, field_name, date)
|
||
self.session.commit()
|
||
else:
|
||
print(f"❌ Record {record_id} not found in __update_date.")
|
||
|
||
def update_date_shorts(self, date: str):
|
||
self.__update_date("date_shorts", date)
|
||
|
||
def update_date_chats(self, date: str):
|
||
self.__update_date("date_videos_chat", date)
|
||
|
||
def get_data(self) -> list[Data]:
|
||
stmt = select(Data)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_youtube_info_for_video(self, video_id: str) -> list[YoutubeInfo]:
|
||
"""Retrieve all YouTubeInfo records associated with a specific TwitchVideo ID."""
|
||
stmt = select(YoutubeInfo).where(YoutubeInfo.twitch_video_id == video_id)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_youtube_info_all(self) -> list[YoutubeInfo]:
|
||
return self.session.scalars(select(YoutubeInfo)).all()
|
||
|
||
def insert(self, data: TwitchVideo | Data | YoutubeInfo | Chapter):
|
||
def insert(self, data: TwitchVideo | Data | YoutubeInfo | Chapter | list[YoutubeInfo] | list[Chapter]):
|
||
"""
|
||
Handles insertion/updating for any of the defined models,
|
||
using session merge/add for SQLAlchemy 2.0 best practices.
|
||
Accepts single objects or lists of objects for collection handling.
|
||
"""
|
||
try:
|
||
if isinstance(data, TwitchVideo):
|
||
# 1. Check if the record exists
|
||
stmt = select(TwitchVideo).where(TwitchVideo.id == data.id)
|
||
existing_record = self.session.scalar(stmt)
|
||
|
||
if existing_record:
|
||
# 2. Update relevant fields on the managed instance
|
||
if existing_record.view_count != data.view_count or existing_record.duration != data.duration:
|
||
existing_record.view_count = data.view_count
|
||
existing_record.duration = data.duration
|
||
print(f"🔄 TwitchVideo record {data.id} updated successfully!")
|
||
# Note: Add more field updates here as needed
|
||
else:
|
||
# 3. Add new instance
|
||
self.session.add(data)
|
||
print(f"➕ New TwitchVideo record with ID {data.id} inserted successfully!")
|
||
elif isinstance(data, Data):
|
||
self.session.merge(data)
|
||
print(f"➕ Data record merged/inserted successfully!")
|
||
elif isinstance(data, YoutubeInfo):
|
||
# Handle single YoutubeInfo object
|
||
self.session.merge(data)
|
||
print(f"➕ YoutubeInfo record merged/inserted successfully!")
|
||
elif isinstance(data, Chapter):
|
||
# Handle single Chapter object
|
||
self.session.merge(data)
|
||
print(f"➕ Chapter record merged/inserted successfully!")
|
||
elif isinstance(data, list):
|
||
# Handle lists of records (YoutubeInfo or Chapter)
|
||
record_type = type(data[0])
|
||
if record_type in (YoutubeInfo, Chapter):
|
||
for item in data:
|
||
self.session.merge(item)
|
||
print(f"➕ Successfully processed {len(data)} records of type {record_type.__name__}.")
|
||
else:
|
||
raise TypeError(f"Unsupported list content type: {record_type.__name__}")
|
||
else:
|
||
raise TypeError(f"Unsupported data type for insertion: {type(data).__name__}")
|
||
|
||
# 4. Commit the transaction
|
||
self.session.flush() # Flush first to ensure IDs are populated
|
||
|
||
except Exception as e:
|
||
self.session.rollback()
|
||
print(f"❌ A critical error occurred during record insertion/update: {e}")
|
||
raise
|
||
|
||
self.session.commit()
|