Fixing critical bug in DB

Dockerfile

@@ -11,12 +11,14 @@ COPY ./requirements-fast.txt ./
 
 COPY ./core ./core
 COPY ./rl_string_helper ./rl_string_helper
+COPY ./database-lib ./database-lib
 
 COPY ./other/sqlite_zstd-0.1.dev1+g5aaeb60-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ./
 
 RUN pip install --no-cache-dir wheel
 
 RUN pip3 install --no-cache-dir ./rl_string_helper
+RUN pip3 install --no-cache-dir ./database-lib
 RUN pip3 install --no-cache-dir ./core
 
 RUN pip3 install --no-cache-dir -r requirements.txt

core/medium_parser/__init__.py

@@ -2,7 +2,7 @@
 from aiohttp_retry import ExponentialRetry
 
 import jinja2
-from .cache_db import SQLiteCacheBackend
+from database_lib import SQLiteCacheBackend
 
 cache = SQLiteCacheBackend('medium_db_cache.sqlite')
 cache.init_db()

core/medium_parser/db_cache_migration.py

@@ -1,31 +0,0 @@
-import sqlite3
-import asyncio
-import pickle
-from cache_db import SQLiteCacheBackend
-
-db_path = "../medium_cache.sqlite"
-
-async def main():
-    conn = sqlite3.connect(db_path)
-    db_cache = SQLiteCacheBackend("medium_db_cache.sqlite")
-    db_cache.init_db()
-
-    c = conn.cursor()
-
-    c.execute("SELECT * FROM responses")
-
-    results = c.fetchall()
-
-    for result in results:
-        value_raw = pickle.loads(result[1])
-        db_cache.push(result[0], await value_raw.text())
-
-    # Close the connections
-    c.close()
-    conn.close()
-
-    db_cache.enable_zstd()
-
-    db_cache.close()
-
-asyncio.run(main())

core/requirements.txt

@@ -1,4 +1,5 @@
 rl_string_helper==0.1.0
+database_lib==0.1.0
 
 loguru==0.6.0
 aiohttp==3.8.5
@@ -7,4 +8,3 @@ tld==0.13
 bs4==0.0.1
 Jinja2==3.1.2
 beautifulsoup4==4.12.2
-# git+https://github.com/phiresky/sqlite-zstd.git#egg=sqlite_zstd&subdirectory=python

database-lib/README.md

@@ -0,0 +1,3 @@
+# Database-Lib
+
+Database-Lib is a library for database operations.

database-lib/database_lib/__init__.py

@@ -0,0 +1 @@
+from .cache_db import SQLiteCacheBackend

database-lib/database_lib/cache_db.py

@@ -6,6 +6,7 @@ from warnings import warn
 try:
     import sqlite_zstd
 except ImportError:
+    logger.debug("Can't use zstd compression. Please install 'sqlite_zstd' package")
     warn("Can't use zstd compression. Please install 'sqlite_zstd' package")
     sqlite_zstd = None
 
@@ -36,6 +37,8 @@ class SQLiteCacheBackend:
         if sqlite_zstd is not None:
             sqlite_zstd.load(self.connection)
 
+        self.migrate_add_index_to_key()
+
     def all(self):
         with self.connection:
             return self.cursor.execute("SELECT * FROM cache").fetchall()
@@ -64,25 +67,53 @@ class SQLiteCacheBackend:
     def init_db(self):
         with self.connection:
             self.cursor.execute("CREATE TABLE IF NOT EXISTS cache (key TEXT PRIMARY KEY, value TEXT)")
+            self.cursor.execute("CREATE INDEX IF NOT EXISTS idx_key ON cache (key)")
 
-    def pull(self, key: str) -> Union[dict, str]:
+    def pull(self, key: str) -> Union[CacheResponse, None]:
         with self.connection:
             cache = self.cursor.execute("SELECT value FROM cache WHERE key = :0", {'0': key}).fetchone()
             if cache:
                 logger.debug("Value found in DB, returning it")
                 return CacheResponse(cache[0])
+            else:
+                logger.debug(f"No value found for key: {key}")
+                return None
 
-    def push(self, key: str, value: str) -> None:
+    def push(self, key: str, value: Union[str, dict]) -> None:
         if isinstance(value, dict):
-            value = json.dumps(value)
+            try:
+                value = json.dumps(value)
+            except TypeError as e:
+                raise ValueError(f"Unable to serialize value to JSON: {e}")
         elif not isinstance(value, str):
-            raise ValueError(f"value argument should be only string type not {type(value).__name__}")
+            raise ValueError(f"value argument should be a string or dict, not {type(value).__name__}")
         with self.connection:
             self.cursor.execute("INSERT OR REPLACE INTO cache VALUES (:0, :1)", {'0': key, '1': value})
-
+            
     def delete(self, key: str) -> None:
         with self.connection:
-            self.cursor.execute("DELETE FROM cache WHERE key = :0", {'0': key})
+            result = self.cursor.execute("SELECT 1 FROM cache WHERE key = :0", {'0': key}).fetchone()
+            if result:
+                self.cursor.execute("DELETE FROM cache WHERE key = :0", {'0': key})
+                logger.debug(f"Deleted key: {key}")
+            else:
+                logger.debug(f"Attempted to delete non-existing key: {key}")
+    
+    def maintenance(self):
+        with self.connection:
+            self.cursor.execute("VACUUM")
+            self.cursor.execute("ANALYZE")
+
+    def migrate_add_index_to_key(self):
+        with self.connection:
+            # Check if the index already exists
+            index_exists = self.cursor.execute("SELECT name FROM sqlite_master WHERE type='index' AND name='idx_key'").fetchone()
+            if not index_exists:
+                # Create the index if it doesn't exist
+                self.cursor.execute("CREATE INDEX idx_key ON cache (key)")
+                logger.info("Index 'idx_key' on column 'key' created successfully.")
+            else:
+                logger.info("Index 'idx_key' on column 'key' already exists.")
 
     def close(self):
         self.__del__()

database-lib/requirements.txt

@@ -0,0 +1,3 @@
+loguru==0.6.0
+
+# git+https://github.com/phiresky/sqlite-zstd.git#egg=sqlite_zstd&subdirectory=python

database-lib/setup.py

@@ -0,0 +1,25 @@
+from setuptools import setup, find_packages
+
+# Function to read the contents of the requirements file
+def read_requirements():
+    with open('requirements.txt', 'r') as req:
+        return req.read().splitlines()
+
+setup(
+    name='database_lib',
+    version='0.1.0',
+    author='Freedium community',
+    author_email='admin@freedium.cfd',
+    description='A dabase helper library from Medium cache',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    url='https://codeberg.org/Freedium-cfd/web',
+    packages=find_packages(),
+    install_requires=read_requirements(),
+    classifiers=[
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: OS Independent',
+    ],
+    python_requires='>=3.7',
+)

database-lib/test.py

@@ -0,0 +1,66 @@
+import unittest
+from database_lib.cache_db import SQLiteCacheBackend
+import os
+
+class TestSQLiteCacheBackend(unittest.TestCase):
+    test_db = 'test_cache.db'
+
+    @classmethod
+    def setUpClass(cls):
+        cls.cache_backend = SQLiteCacheBackend(cls.test_db)
+        cls.cache_backend.init_db()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.cache_backend.close()
+        os.remove(cls.test_db)
+
+    def test_push_and_pull(self):
+        key, value = 'test_key', {"Hii": "Hii"}
+        self.cache_backend.push(key, value)
+        result = self.cache_backend.pull(key)
+        self.assertEqual(result.json(), value, "The pulled value should match the pushed value.")
+
+    def test_delete(self):
+        key, value = 'delete_key', 'delete_value'
+        self.cache_backend.push(key, value)
+        self.cache_backend.delete(key)
+        result = self.cache_backend.pull(key)
+        self.assertIsNone(result, "The result should be None after deletion.")
+
+    def test_all_length(self):
+        initial_length = self.cache_backend.all_length()
+        self.cache_backend.push('length_key', 'length_value')
+        new_length = self.cache_backend.all_length()
+        self.assertEqual(new_length, initial_length + 1, "The length should increase by 1 after adding a new item.")
+
+    def test_random(self):
+        # Ensure there is at least one item
+        self.cache_backend.push('random_key', 'random_value')
+        result = self.cache_backend.random(1)
+        self.assertTrue(len(result) > 0, "Should return at least one item.")
+
+    def test_migration_add_index_to_key(self):
+        # Forcefully remove the index if it exists to simulate a scenario where the migration is needed.
+        self.cache_backend.cursor.execute("DROP INDEX IF EXISTS idx_key")
+        self.cache_backend.connection.commit()
+
+        # Call the migration method to add the index.
+        self.cache_backend.migrate_add_index_to_key()
+
+        # Verify the index has been created.
+        index_exists = self.cache_backend.cursor.execute("SELECT name FROM sqlite_master WHERE type='index' AND name='idx_key'").fetchone()
+        self.assertIsNotNone(index_exists, "The index 'idx_key' should exist after migration.")
+
+        # Call the migration method again to simulate the scenario where the index already exists.
+        self.cache_backend.migrate_add_index_to_key()
+
+        # Verify that the index still exists and there are no errors.
+        index_exists = self.cache_backend.cursor.execute("SELECT name FROM sqlite_master WHERE type='index' AND name='idx_key'").fetchone()
+        self.assertIsNotNone(index_exists, "The index 'idx_key' should still exist after calling migration again.")
+
+        # self.cache_backend.cursor.execute("DROP INDEX IF EXISTS idx_key")
+        # self.cache_backend.connection.commit()
+
+if __name__ == '__main__':
+    unittest.main()