about summary refs log tree commit diff
path: root/.venv/lib/python3.12/site-packages/core/providers/database/tokens.py
diff options
context:
space:
mode:
authorS. Solomon Darnell2025-03-28 21:52:21 -0500
committerS. Solomon Darnell2025-03-28 21:52:21 -0500
commit4a52a71956a8d46fcb7294ac71734504bb09bcc2 (patch)
treeee3dc5af3b6313e921cd920906356f5d4febc4ed /.venv/lib/python3.12/site-packages/core/providers/database/tokens.py
parentcc961e04ba734dd72309fb548a2f97d67d578813 (diff)
downloadgn-ai-master.tar.gz
two version of R2R are here HEAD master
Diffstat (limited to '.venv/lib/python3.12/site-packages/core/providers/database/tokens.py')
-rw-r--r--.venv/lib/python3.12/site-packages/core/providers/database/tokens.py67
1 files changed, 67 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/core/providers/database/tokens.py b/.venv/lib/python3.12/site-packages/core/providers/database/tokens.py
new file mode 100644
index 00000000..7d30c326
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/core/providers/database/tokens.py
@@ -0,0 +1,67 @@
+from datetime import datetime, timedelta
+from typing import Optional
+
+from core.base import Handler
+
+from .base import PostgresConnectionManager
+
+
+class PostgresTokensHandler(Handler):
+    TABLE_NAME = "blacklisted_tokens"
+
+    def __init__(
+        self, project_name: str, connection_manager: PostgresConnectionManager
+    ):
+        super().__init__(project_name, connection_manager)
+
+    async def create_tables(self):
+        query = f"""
+        CREATE TABLE IF NOT EXISTS {self._get_table_name(PostgresTokensHandler.TABLE_NAME)} (
+            id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+            token TEXT NOT NULL,
+            blacklisted_at TIMESTAMPTZ DEFAULT NOW()
+        );
+        CREATE INDEX IF NOT EXISTS idx_{self.project_name}_{PostgresTokensHandler.TABLE_NAME}_token
+        ON {self._get_table_name(PostgresTokensHandler.TABLE_NAME)} (token);
+        CREATE INDEX IF NOT EXISTS idx_{self.project_name}_{PostgresTokensHandler.TABLE_NAME}_blacklisted_at
+        ON {self._get_table_name(PostgresTokensHandler.TABLE_NAME)} (blacklisted_at);
+        """
+        await self.connection_manager.execute_query(query)
+
+    async def blacklist_token(
+        self, token: str, current_time: Optional[datetime] = None
+    ):
+        if current_time is None:
+            current_time = datetime.utcnow()
+
+        query = f"""
+        INSERT INTO {self._get_table_name(PostgresTokensHandler.TABLE_NAME)} (token, blacklisted_at)
+        VALUES ($1, $2)
+        """
+        await self.connection_manager.execute_query(
+            query, [token, current_time]
+        )
+
+    async def is_token_blacklisted(self, token: str) -> bool:
+        query = f"""
+        SELECT 1 FROM {self._get_table_name(PostgresTokensHandler.TABLE_NAME)}
+        WHERE token = $1
+        LIMIT 1
+        """
+        result = await self.connection_manager.fetchrow_query(query, [token])
+        return bool(result)
+
+    async def clean_expired_blacklisted_tokens(
+        self,
+        max_age_hours: int = 7 * 24,
+        current_time: Optional[datetime] = None,
+    ):
+        if current_time is None:
+            current_time = datetime.utcnow()
+        expiry_time = current_time - timedelta(hours=max_age_hours)
+
+        query = f"""
+        DELETE FROM {self._get_table_name(PostgresTokensHandler.TABLE_NAME)}
+        WHERE blacklisted_at < $1
+        """
+        await self.connection_manager.execute_query(query, [expiry_time])