Update Akeno/utils/base_sqlite.py

Akeno/utils/base_sqlite.py

@@ -1,34 +1,43 @@
-import sqlitecloud
-import asyncio
-from config import SQLITE_CLOUD_URL
-
-conn = sqlitecloud.connect(SQLITE_CLOUD_URL)
-cursor = conn.cursor()
-cursor.execute('USE DATABASE akeno-dev')
-
-cursor.execute('''
-    CREATE TABLE IF NOT EXISTS prefixes (
-        user_id INTEGER PRIMARY KEY,
-        prefix TEXT NOT NULL
-    )
-''')
-conn.commit()
+from sqlalchemy import create_engine, Column, Integer, String
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker, scoped_session
 
-async def set_prefix_in_db(user_id: int, prefix: str):
-    cursor.execute('''
-        INSERT INTO prefixes (user_id, prefix)
-        VALUES (?, ?)
-        ON CONFLICT(user_id) DO UPDATE SET prefix=excluded.prefix
-    ''', (user_id, prefix))
-    conn.commit()
+from config import DB_URI
 
-async def get_prefix(user_id: int):
-    cursor.execute('SELECT prefix FROM prefixes WHERE user_id=?', (user_id,))
-    result = cursor.fetchone()
+if DB_URI and DB_URI.startswith("postgres://"):
+    DB_URI = DB_URI.replace("postgres://", "postgresql://", 1)
+
+engine = create_engine(DB_URI)
+SessionLocal = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=engine))
+
+Base = declarative_base()
+
+# Define the Prefix model
+class Prefix(Base):
+    __tablename__ = "prefixes"
 
-    if result:
-        return result[0]
-    else:
-        return None
+    user_id = Column(Integer, primary_key=True, index=True)
+    prefix = Column(String, nullable=False)
 
-conn.close()
+Base.metadata.create_all(bind=engine)
+
+async def set_prefix_in_db(user_id: int, prefix: str):
+    db = SessionLocal()
+    try:
+        user_prefix = db.query(Prefix).filter(Prefix.user_id == user_id).first()
+        if user_prefix:
+            user_prefix.prefix = prefix
+        else:
+            new_prefix = Prefix(user_id=user_id, prefix=prefix)
+            db.add(new_prefix)
+        db.commit()
+    finally:
+        db.close()
+
+async def get_prefix(user_id: int):
+    db = SessionLocal()
+    try:
+        user_prefix = db.query(Prefix).filter(Prefix.user_id == user_id).first()
+        return user_prefix.prefix if user_prefix else None
+    finally:
+        db.close()