From 6a9c1038d22d8eff35088d842b36f13e937613d2 Mon Sep 17 00:00:00 2001
From: yuki <yvki@riseup.net>
Date: Tue, 7 Oct 2025 03:09:25 -0300
Subject: [PATCH] add sqlite database

---
 .gitignore        |  2 ++
 bot.py            | 36 ++++++++++++++++++++++++++++++++++--
 commands/utils.py |  2 --
 main.py           |  5 ++---
 4 files changed, 38 insertions(+), 7 deletions(-)
 delete mode 100644 commands/utils.py

diff --git a/.gitignore b/.gitignore
index b9a7584..4fb2bcd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,8 @@ config.toml
 
 .vscode
 
+database.db
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[codz]
diff --git a/bot.py b/bot.py
index df6b610..ee70bff 100644
--- a/bot.py
+++ b/bot.py
@@ -1,17 +1,23 @@
 import os
 import importlib
+import aiosqlite
 
 import nerimity
 import pylast
 
 class Bot(nerimity.Client):
     """Extended client class for extra functionality."""
-    def __init__(self, lastfm_api_key: str, lastfm_api_secret: str, *args, **kwargs):
+    def __init__(self, lastfm_api_key: str, lastfm_api_secret: str, db_file: str = "database.db", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # initialize last.fm client
         self.lastfm = pylast.LastFMNetwork(
             api_key = lastfm_api_key,
             api_secret = lastfm_api_secret
         )
-        super().__init__(*args, **kwargs)
+
+        self.db = None
+        self.db_file = db_file
 
     def load_commands(self, commands_dir: str):
         for filename in os.listdir(commands_dir):
@@ -22,3 +28,29 @@ class Bot(nerimity.Client):
                     module.setup(self)
                 print(f"Loaded {module_name}")
         print("Registered commands:", self.commands)
+
+    async def init_db(self):
+        print("Initializing database...")
+        self.db = await aiosqlite.connect(self.db_file)
+        await self.db.execute("""
+            CREATE TABLE IF NOT EXISTS lastfm_users (
+                user_id TEXT PRIMARY KEY,
+                lastfm_username TEXT NOT NULL
+            )
+        """)
+        await self.db.commit()
+
+    async def set_lastfm(self, user_id: str, username: str):
+        await self.db.execute(
+            "INSERT OR REPLACE INTO lastfm_users (user_id, lastfm_username) VALUES (?, ?)", 
+            (user_id, username)
+        )
+        await self.db.commit()
+
+    async def get_lastfm(self, user_id: str):
+        async with self.db.execute(
+            "SELECT lastfm_username FROM lastfm_users WHERE user_id = ?",
+            (user_id,)
+        ) as cursor:
+            row = await cursor.fetchone()
+            return row[0] if row else None
diff --git a/commands/utils.py b/commands/utils.py
deleted file mode 100644
index b8ccc6e..0000000
--- a/commands/utils.py
+++ /dev/null
@@ -1,2 +0,0 @@
-def error_msg(message: str):
-    return f"[#e5323b][Error] [#reset]{message}"
diff --git a/main.py b/main.py
index 439b494..7377326 100644
--- a/main.py
+++ b/main.py
@@ -1,7 +1,5 @@
 import tomllib
-
 import nerimity
-import pylast
 
 import bot
 
@@ -10,8 +8,8 @@ with open("config.toml", "rb") as f:
 
 # nerimity client
 bot = bot.Bot(
-    token = config['token'],
     prefix = '!',
+    token = config['token'],
     lastfm_api_key = config['api_key'],
     lastfm_api_secret = config['api_secret']
 )
@@ -21,5 +19,6 @@ bot.load_commands("commands")
 @bot.listen("on_ready")
 async def on_ready(params):
     print(f"Logged in as {bot.account.username}")
+    await bot.init_db()
 
 bot.run()