From b9e95d4ed06c094ce7d5e0bf4708107d467b7d25 Mon Sep 17 00:00:00 2001
From: Vlad Pronsky <v.pronsky@gmail.com>
Date: Thu, 18 Apr 2024 05:12:10 +0300
Subject: [PATCH] simplify custom poll accounts control (#138)

---
 twscrape/accounts_pool.py | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/twscrape/accounts_pool.py b/twscrape/accounts_pool.py
index 79114f7..c41f5be 100644
--- a/twscrape/accounts_pool.py
+++ b/twscrape/accounts_pool.py
@@ -238,24 +238,16 @@ class AccountsPool:
         """
         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 {self._order_by}
-        LIMIT 1
-        """
+    async def _get_and_lock(self, queue: str, condition: str):
+        # if space in condition, it's a subquery, otherwise it's username
+        condition = f"({condition})" if " " in condition else f"'{condition}'"
 
         if int(sqlite3.sqlite_version_info[1]) >= 35:
             qs = f"""
             UPDATE accounts SET
                 locks = json_set(locks, '$.{queue}', datetime('now', '+15 minutes')),
                 last_used = datetime({utc.ts()}, 'unixepoch')
-            WHERE username = ({q1})
+            WHERE username = {condition}
             RETURNING *
             """
             rs = await fetchone(self._db_file, qs)
@@ -266,7 +258,7 @@ class AccountsPool:
                 locks = json_set(locks, '$.{queue}', datetime('now', '+15 minutes')),
                 last_used = datetime({utc.ts()}, 'unixepoch'),
                 _tx = '{tx}'
-            WHERE username = ({q1})
+            WHERE username = {condition}
             """
             await execute(self._db_file, qs)
 
@@ -275,6 +267,20 @@ class AccountsPool:
 
         return Account.from_rs(rs) if rs else None
 
+    async def get_for_queue(self, queue: str):
+        q = 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 {self._order_by}
+        LIMIT 1
+        """
+
+        return await self._get_and_lock(queue, q)
+
     async def get_for_queue_or_wait(self, queue: str) -> Account | None:
         msg_shown = False
         while True: