208 lines
8.7 KiB
Python
208 lines
8.7 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)
|
||
|
||
# Relationship back to YoutubeInfo (uselist=False enforces a 1-to-1 structure)
|
||
youtube_info = relationship("YoutubeInfo", back_populates="twitch_video", uselist=False)
|
||
|
||
# 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_chat = Column(String, nullable=False)
|
||
|
||
class YoutubeInfo(Base):
|
||
__tablename__ = "youtube_info"
|
||
|
||
id = Column(String, ForeignKey("twitch_videos.id"), primary_key=True)
|
||
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 Database:
|
||
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 __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 get_data(self) -> list[Data]:
|
||
stmt = select(Data)
|
||
return self.session.scalars(stmt).all()
|
||
|
||
def get_youtube_info(self, record_id) -> YoutubeInfo | None:
|
||
return self.session.get(YoutubeInfo, record_id)
|
||
|
||
def insert(self, data: TwitchVideo | Data | YoutubeInfo):
|
||
# SQLAlchemy 2.0 context manager handles commit/rollback automatically
|
||
try:
|
||
if isinstance(data, TwitchVideo):
|
||
# 1. Query for the existing row using SQLAlchemy 2.0 scalar execution
|
||
stmt = select(TwitchVideo).where(TwitchVideo.id == data.id)
|
||
existing_record = self.session.scalar(stmt)
|
||
|
||
if existing_record:
|
||
if existing_record.view_count != data.view_count or existing_record.duration != data.duration:
|
||
# 2. Update the columns on the managed instance
|
||
existing_record.view_count = data.view_count
|
||
existing_record.duration = data.duration
|
||
print(f"🔄 Record {data.id} updated successfully!")
|
||
else:
|
||
# 3. Add new instance if it doesn't exist
|
||
self.session.add(data)
|
||
print("➕ New TwitchVideo record inserted successfully!")
|
||
else:
|
||
self.session.merge(data)
|
||
print(f"➕ {type(data).__name__} record inserted successfully!")
|
||
# 4. Flush changes and commit the transaction
|
||
self.session.commit()
|
||
|
||
except Exception as e:
|
||
self.session.rollback()
|
||
print(f"❌ An error occurred, transaction rolled back: {e}")
|