twscrape/twapi/accounts_pool.py

# ruff: noqa: E501
import asyncio

from fake_useragent import UserAgent

from .account import Account
from .db import add_init_query, execute, fetchall, fetchone
from .logger import logger
from .login import login


class AccountsPool:
    def __init__(self, db_file="accounts.db"):
        self._db_file = db_file
        add_init_query(db_file, Account.create_sql())

    async def add_account(
        self,
        username: str,
        password: str,
        email: str,
        email_password: str,
        user_agent: str | None = None,
        proxy: str | None = None,
    ):
        qs = "SELECT * FROM accounts WHERE username = :username"
        rs = await fetchone(self._db_file, qs, {"username": username})
        if rs:
            return

        account = Account(
            username=username,
            password=password,
            email=email,
            email_password=email_password,
            user_agent=user_agent or UserAgent().safari,
            active=False,
            locks={},
            headers={},
            cookies={},
            proxy=proxy,
        )
        await self.save(account)

    async def get(self, username: str):
        qs = "SELECT * FROM accounts WHERE username = :username"
        rs = await fetchone(self._db_file, qs, {"username": username})
        if not rs:
            raise ValueError(f"Account {username} not found")
        return Account.from_rs(rs)

    async def get_all(self):
        qs = "SELECT * FROM accounts"
        rs = await fetchall(self._db_file, qs)
        return [Account.from_rs(x) for x in rs]

    async def save(self, account: Account):
        data = account.to_rs()
        cols = list(data.keys())

        qs = f"""
        INSERT INTO accounts ({",".join(cols)}) VALUES ({",".join([f":{x}" for x in cols])})
        ON CONFLICT DO UPDATE SET {",".join([f"{x}=excluded.{x}" for x in cols])}
        """
        await execute(self._db_file, qs, data)

    async def login(self, account: Account):
        try:
            await login(account)
        except Exception as e:
            logger.error(f"Error logging in to {account.username}: {e}")
        finally:
            await self.save(account)

    async def login_all(self):
        qs = "SELECT * FROM accounts WHERE active = false AND error_msg IS NULL"
        rs = await fetchall(self._db_file, qs)

        accounts = [Account.from_rs(rs) for rs in rs]
        # await asyncio.gather(*[login(x) for x in self.accounts])

        for i, x in enumerate(accounts, start=1):
            logger.info(f"[{i}/{len(accounts)}] Logging in {x.username} - {x.email}")
            await self.login(x)

    async def set_active(self, username: str, active: bool):
        qs = "UPDATE accounts SET active = :active WHERE username = :username"
        await execute(self._db_file, qs, {"username": username, "active": active})

    async def lock_until(self, username: str, queue: str, unlock_at: int):
        # unlock_at is unix timestamp
        qs = f"""
        UPDATE accounts SET locks = json_set(locks, '$.{queue}', datetime({unlock_at}, 'unixepoch'))
        WHERE username = :username
        """
        await execute(self._db_file, qs, {"username": username})

    async def unlock(self, username: str, queue: str):
        qs = f"""
        UPDATE accounts SET locks = json_remove(locks, '$.{queue}')
        WHERE username = :username
        """
        await execute(self._db_file, qs, {"username": username})

    async def get_for_queue(self, queue: str):
        q1 = f"""
        SELECT username FROM accounts
        WHERE active = true AND (
            locks IS NULL
            OR json_extract(locks, '$.{queue}') IS NULL
            OR json_extract(locks, '$.{queue}') < datetime('now')
        )
        ORDER BY RANDOM()
        LIMIT 1
        """

        q2 = f"""
        UPDATE accounts SET locks = json_set(locks, '$.{queue}', datetime('now', '+15 minutes'))
        WHERE username = ({q1})
        RETURNING *
        """

        rs = await fetchone(self._db_file, q2)
        return Account.from_rs(rs) if rs else None

    async def get_for_queue_or_wait(self, queue: str) -> Account:
        while True:
            account = await self.get_for_queue(queue)
            if not account:
                logger.debug(f"No accounts available for queue '{queue}' (sleeping for 5 sec)")
                await asyncio.sleep(5)
                continue

            logger.debug(f"Using account {account.username} for queue '{queue}'")
            return account

    async def stats(self):
        def by_queue(queue: str):
            return f"""
            SELECT COUNT(*) FROM accounts
            WHERE json_extract(locks, '$.{queue}') IS NOT NULL AND json_extract(locks, '$.{queue}') > datetime('now')
            """

        config = [
            ("total", "SELECT COUNT(*) FROM accounts"),
            ("active", "SELECT COUNT(*) FROM accounts WHERE active = true"),
            ("inactive", "SELECT COUNT(*) FROM accounts WHERE active = false"),
            ("locked_search", by_queue("search")),
        ]

        qs = f"SELECT {','.join([f'({q}) as {k}' for k, q in config])}"
        rs = await fetchone(self._db_file, qs)
        return dict(rs) if rs else {}