From 36cbd4c848c7c6fea2f8400c45ecad55e15461af Mon Sep 17 00:00:00 2001
From: nsde <neostyxde@gmail.com>
Date: Fri, 21 Jul 2023 21:54:53 +0200
Subject: [PATCH] Added simple database demo

---
 .gitignore    |   6 +++
 api/closed.py | 139 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 145 insertions(+)
 create mode 100644 api/closed.py

diff --git a/.gitignore b/.gitignore
index 68bc17f..0144ec4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,9 @@
+*.db
+*.sqlite3
+*.sql
+*.mysql
+*.sqlite
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/api/closed.py b/api/closed.py
new file mode 100644
index 0000000..1604283
--- /dev/null
+++ b/api/closed.py
@@ -0,0 +1,139 @@
+import os
+import sys
+import uuid
+import time
+import asyncio
+import aiosqlite
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+db_uri = os.getenv('DB_PATH')
+
+try:
+    os.remove(db_uri)
+except FileNotFoundError:
+    pass
+
+async def to_dict(cursor, row):
+    return {col[0]: row[idx] for idx, col in enumerate(cursor.description)}
+
+async def connect_db():
+    """Creates a connection to the database"""
+
+    return await aiosqlite.connect(db_uri)
+
+async def prepare() -> None:
+    """Creates the database tables"""
+
+    keys_db = await connect_db()
+    await keys_db.execute(
+        """
+        CREATE TABLE IF NOT EXISTS closed_keys (
+            id TEXT PRIMARY KEY,
+            key TEXT,
+            source TEXT DEFAULT 'unknown',
+            created_at INTEGER,
+            last_used INTEGER,
+            uses_count INTEGER DEFAULT 0,
+            tokens_generated INTEGER DEFAULT 0,
+            active BOOLEAN,
+            working BOOLEAN,
+            tags TEXT DEFAULT ''
+        )
+        """
+    )
+    await keys_db.commit()
+
+async def add_key(
+    key: str,
+    source: str='unknown',
+    tags: list=None
+) -> None:
+    """Adds a new key to the database"""
+
+    tags = tags or []
+
+    keys_db = await connect_db()
+
+    parameters = {
+        "id": str(uuid.uuid4()),
+        "key": key,
+        "source": source,
+        "created_at": int(time.time()),
+        "last_used": -1,
+        "uses_count": 0,
+        "tokens_generated": 0,
+        "active": True,
+        "working": True,
+        "tags": '/'.join(tags),
+    }
+
+    sep = ', '
+    query = f"""
+        INSERT INTO closed_keys ({sep.join(parameters.keys())})
+        VALUES ({sep.join([f':{key}' for key in parameters.keys()])})
+    """
+
+    await keys_db.execute(query, parameters)
+    await keys_db.commit()
+
+async def get_working_key() -> dict:
+    """Returns a working key"""
+
+    keys_db = await connect_db()
+
+    async with keys_db.execute('SELECT * FROM closed_keys WHERE working = 1') as cursor:
+        async for row in cursor:
+            return await to_dict(cursor, row)
+
+    return None
+
+asyncio.run(prepare())
+
+async def key_stopped_working(key: str) -> None:
+    """Marks a key as stopped working"""
+
+    keys_db = await connect_db()
+
+    await keys_db.execute(
+        """
+        UPDATE closed_keys
+        SET working = 0
+        WHERE key = :key
+        """,
+        {'key': key}
+    )
+    await keys_db.commit()
+
+async def key_was_used(key: str, num_tokens: int) -> None:
+    """Updates the stats of a key"""
+
+    keys_db = await connect_db()
+
+    # set last_used to int of time.time(), adds one to uses_count and adds num_tokens to tokens_generated
+    await keys_db.execute(
+        """
+        UPDATE closed_keys
+        SET last_used = :last_used, uses_count = uses_count + 1, tokens_generated = tokens_generated + :tokens_generated
+        WHERE key = :key
+        """,
+        {
+            'key': key,
+            'last_used': int(time.time()),
+            'tokens_generated': num_tokens
+        }
+    )
+    await keys_db.commit()
+
+async def demo():
+    await add_key('sk-non-working-key')
+    await key_stopped_working('sk-non-working-key')
+    await add_key('sk-working-key')
+    key = await get_working_key()
+    print(key)
+
+if __name__ == '__main__':
+    asyncio.run(demo())
+    os.system(f'pkill -f {os.path.basename(__file__)}')