kemono2/src/lib/artist.py

import datetime
from typing import Optional, TypedDict, Literal, Dict

from src.config import Configuration
from src.internals.cache.redis import get_conn
from src.internals.database.database import cached_query, get_cursor, get_pool
from src.internals.serializers import unsafe_dumper, unsafe_loader
from src.internals.serializers.artist import (
    deserialize_artist,
    deserialize_artists,
    serialize_artist,
    serialize_artists,
)
from src.utils.utils import clear_web_cache_for_creator_links

class TDArtist(TypedDict):
    id: str
    name: str
    service: str
    indexed: str
    updated: str
    public_id: str
    relation_id: int


class TDArtistWithFavs(TDArtist):
    count: int


def get_top_artists_by_faves(offset, count, reload=False) -> list[TDArtistWithFavs]:
    key = f"top_artists:{offset}:{count}"
    params = dict(offset=offset, limit=count)
    query = """
        SELECT
            artists.id,
            artists.name,
            artists.service,
            artists.indexed,
            artists.updated,
            artists.public_id,
            artists.relation_id,
            fc.favorite_count AS count
        FROM
            lookup AS artists
            INNER JOIN
            favorite_counts AS fc
            ON
                artists.id = fc.artist_id
                AND
                artists.service = fc.service
        WHERE
            (artists.id, artists.service) NOT IN (
                SELECT
                    id,
                    service
                FROM
                    dnp
            )
        ORDER BY
            count DESC
        OFFSET %(offset)s
        LIMIT %(limit)s
    """

    return cached_query(
        query,
        key,
        params,
        serialize_artists,
        deserialize_artists,
        reload,
        ex=int(datetime.timedelta(hours=24).total_seconds()),
        lock_enabled=True,
    )


def get_random_artist_keys(count, reload=False):
    key = f"random_artist_keys:{count}"
    query = "SELECT id, service FROM lookup ORDER BY random() LIMIT %s"
    return cached_query(query, key, (count,), unsafe_dumper, unsafe_loader, reload, lock_enabled=True)


def get_artist(service: str, artist_id: str, reload: bool = False) -> TDArtist:
    key = f"artist:{service}:{artist_id}"
    params = dict(artist_id=artist_id, service=service)
    id_filter = (
        "(id = %(artist_id)s or public_id = %(artist_id)s)"
        if service in ("onlyfans", "fansly", "candfans", "patreon", "fanbox", "boosty")
        else "id = %(artist_id)s"
    )
    query = f"""
        SELECT
            id,
            name,
            service,
            indexed,
            updated,
            public_id,
            relation_id
        FROM
            lookup
        WHERE
            {id_filter}
            AND service = %(service)s
            AND (id, service) NOT IN (
                SELECT
                    id,
                    service
                FROM dnp
            );
    """
    return cached_query(
        query,
        key,
        params,
        serialize_artist,
        deserialize_artist,
        reload,
        True,
        86400,
        Configuration().redis["default_ttl"],
    )


def get_artists_by_update_time(offset, limit, reload=False) -> list[TDArtist]:
    key = f"artists_by_update_time:{offset}:{limit}"
    params = dict(offset=offset, limit=limit)
    query = """
        SELECT
            id,
            name,
            service,
            indexed,
            updated,
            public_id,
            relation_id
        FROM
            lookup
        WHERE
            (id, service) NOT IN (
                SELECT
                    id,
                    service
                FROM
                    dnp
            )
        ORDER BY
            updated DESC
        OFFSET %(offset)s
        LIMIT %(limit)s
    """

    return cached_query(query, key, params, serialize_artists, deserialize_artists, reload)


def get_fancards_by_artist(artist_id, reload=False):
    key = f"fancards:{artist_id}"
    query = "select * from fanbox_fancards left join files on file_id = files.id where user_id = %s and file_id is not null order by added desc"
    return cached_query(query, key, (artist_id,), reload=reload)


def create_unapproved_link_request(from_artist, to_artist, user_id, reason: Optional[str]):
    query = """
    INSERT INTO unapproved_link_requests (from_service, from_id, to_service, to_id, requester_id, reason)
    VALUES (%s, %s, %s, %s, %s, %s)
    ON CONFLICT DO NOTHING
    """
    cur = get_cursor()
    cur.execute(
        query,
        (from_artist["service"], from_artist["id"], to_artist["service"], to_artist["id"], user_id, reason or None),
    )


class TDUnapprovedLink(TypedDict):
    id: int
    from_service: str
    from_id: str
    to_service: str
    to_id: str
    reason: str
    requester_id: int
    status: Literal["pending", "approved", "rejected"]
    from_creator: Dict
    to_creator: Dict
    requester: Dict


def get_unapproved_links_with_artists() -> list[TDUnapprovedLink]:
    query = """
    SELECT
        unapproved_link_requests.* ,
        row_to_json(
            from_creator.*
        ) AS from_creator,
        row_to_json(
            to_creator.*
        ) AS to_creator,
        row_to_json(
            requester.*
        ) AS requester
    FROM
        unapproved_link_requests
        JOIN
        lookup AS from_creator
        ON
            from_service = from_creator.service
            AND
            from_id = from_creator.id
        JOIN
        lookup AS to_creator
        ON
            to_service = to_creator.service
            AND
            to_id = to_creator.id
        JOIN
        account AS requester
        ON
            requester_id = requester.id
    WHERE
        status = 'pending'
    ORDER BY
        unapproved_link_requests.id ASC
    """
    cur = get_cursor()
    cur.execute(query)
    return cur.fetchall()


def reject_unapproved_link_request(request_id: int):
    query = """
    UPDATE unapproved_link_requests
    SET status = 'rejected'
    WHERE id = %s
    """
    cur = get_cursor()
    cur.execute(query, (request_id,))


def approve_unapproved_link_request(request_id: int):
    # move from the not same to the diffrent one
    query1 = """
    UPDATE unapproved_link_requests
    SET status = 'approved'
    WHERE id = %s
     OR (from_service, to_service, from_id, to_id) in (SELECT from_service, to_service, from_id, to_id from unapproved_link_requests WHERE id = %s)
     OR (from_service, to_service, from_id, to_id) in (SELECT to_service, from_service, to_id, from_id from unapproved_link_requests WHERE id = %s)
    Returning from_service, to_service, from_id, to_id
    """
    query2 = """
    WITH next_relation_table AS (
        SELECT nextval('lookup_relation_id_seq') as next_id
    )
    UPDATE lookup creators
    SET relation_id = coalesce(from_creator.relation_id, to_creator.relation_id, (SELECT next_id FROM next_relation_table))
    FROM lookup from_creator, lookup to_creator, unapproved_link_requests link
    WHERE
        (
            (creators.id = link.from_id and creators.service = link.from_service)
            OR
            (creators.id = link.to_id and creators.service = link.to_service)
        ) AND
        link.id = %s AND
        from_creator.service = link.from_service AND
        from_creator.id = link.from_id AND
        to_creator.service = link.to_service AND
        to_creator.id = link.to_id
    """
    cursor = get_cursor()
    cursor.execute("BEGIN;")
    cursor.execute(query1, (request_id, request_id, request_id))
    update_result = cursor.fetchone()
    if not update_result:
        raise Exception("Failed to change status of request")
    cursor.execute(query2, (request_id,))
    cursor.execute("COMMIT;")
    redis = get_conn()
    redis.delete(f"linked_accounts:{update_result['from_service']}:{update_result['from_id']}")
    redis.delete(f"linked_accounts:{update_result['to_service']}:{update_result['to_id']}")
    redis.delete(f"artist:{update_result['from_service']}:{update_result['from_id']}")
    redis.delete(f"artist:{update_result['to_service']}:{update_result['to_id']}")
    clear_web_cache_for_creator_links(update_result["from_service"], update_result["from_id"])
    clear_web_cache_for_creator_links(update_result["to_service"], update_result["to_id"])


def delete_creator_link(service: str, creator_id: str):
    query = """
    UPDATE lookup
    SET relation_id = NULL
    WHERE service = %s AND id = %s
    """
    cur = get_cursor()
    cur.execute(query, (service, creator_id))


def get_linked_creators(service: str, creator_id: str):
    key = f"linked_accounts:{service}:{creator_id}"
    query = """
    SELECT * FROM lookup l1
    JOIN lookup l2 ON l1.relation_id = l2.relation_id
    WHERE l1.service = %s AND l1.id = %s
    AND l1.ctid != l2.ctid
    ORDER BY l1.id, l1.service, l2.id, l2.service
    """
    return cached_query(query, key, (service, creator_id), reload=True)


def nload_query_artists(artist_ids: list[tuple[str, str]], dict_result=False):
    # todo remake this like the nload for posts
    if not artist_ids:
        return []
    keys = [f"artist:{service}:{artist_id}" for service, artist_id in artist_ids]
    redis = get_conn()
    cached = (deserialize_artist(artist) for artist in (artist_str for artist_str in redis.mget(keys) if artist_str))
    in_cache = {(artist["service"], artist["id"]): artist for artist in cached if artist}
    for service, artist_id in artist_ids:
        if in_cache.get((service, artist_id)) is None:
            in_cache[(service, artist_id)] = get_artist(service, artist_id) or {}
    if dict_result:
        return in_cache

    artists = []
    for service, artist_id in artist_ids:
        artists.append(in_cache.get((service, artist_id)))
    return artists