diff --git a/src/invidious_exporter/__init__.py b/src/invidious_exporter/__init__.py index 17fc057..7b49000 100644 --- a/src/invidious_exporter/__init__.py +++ b/src/invidious_exporter/__init__.py @@ -2,9 +2,17 @@ import psycopg from argparse import ArgumentParser from os import makedirs, path +from asyncio import run from . import database +async def connect_and_export(connection_string: str, username: str): + async with await psycopg.AsyncConnection.connect(connection_string, row_factory = psycopg.rows.namedtuple_row) as conn: + async with conn.cursor() as cur: + data_export = await database.export_user(cur, username) + + print(data_export.to_dict(), end='') + def main(): parser = ArgumentParser( prog='invidious-exporter', @@ -15,8 +23,6 @@ def main(): args = parser.parse_args() - with psycopg.connect(args.database, row_factory = psycopg.rows.namedtuple_row) as conn: - with conn.cursor() as cur: - data_export = database.export_user(cur, args.username) + run(connect_and_export(args.database, args.username)) - print(data_export.to_dict(), end='') \ No newline at end of file + \ No newline at end of file diff --git a/src/invidious_exporter/database.py b/src/invidious_exporter/database.py index 40ac2df..49704df 100644 --- a/src/invidious_exporter/database.py +++ b/src/invidious_exporter/database.py @@ -4,7 +4,7 @@ from .objects import Channel, Video, Playlist, PlaylistPrivacy, InvidiousJsonDat import json -def get_playlists(cur: Cursor, user_email: str) -> list[Playlist]: +async def get_playlists(cur: Cursor, user_email: str) -> list[Playlist]: sql = """ SELECT p.title, @@ -21,7 +21,7 @@ def get_playlists(cur: Cursor, user_email: str) -> list[Playlist]: p.index, p.title, p.description, p.privacy; """ - cur.execute(sql, (user_email,)) + await cur.execute(sql, (user_email,)) return [ Playlist( @@ -30,18 +30,18 @@ def get_playlists(cur: Cursor, user_email: str) -> list[Playlist]: privacy=PlaylistPrivacy[row.privacy], video_ids=row.video_ids or [] ) - for row in cur.fetchall() + for row in await cur.fetchall() ] -def export_user(cur: Cursor, user_email: str) -> InvidiousJsonDataExport: - cur.execute("SELECT subscriptions, watched, preferences FROM users WHERE email = %s", (user_email,)) +async def export_user(cur: Cursor, user_email: str) -> InvidiousJsonDataExport: + await cur.execute("SELECT subscriptions, watched, preferences FROM users WHERE email = %s", (user_email,)) - data = cur.fetchone() + data = await cur.fetchone() subscribed_channels_ids = data.subscriptions watched_videos_ids = data.watched user_preferences = json.loads(data.preferences) - playlists = get_playlists(cur, user_email) + playlists = await get_playlists(cur, user_email) data_export = InvidiousJsonDataExport( subscriptions = [Channel(id=c_id) for c_id in subscribed_channels_ids],