From 14bbf1c1c52c184c1c18e21b1fa9041043b4a4b6 Mon Sep 17 00:00:00 2001 From: ZhymabekRoman Date: Sat, 21 Sep 2024 16:18:22 +0500 Subject: [PATCH] db cache: fixing broken json parsing --- database-lib/database_lib/main.py | 106 +++++++++++++++++++++------- medium-parser/medium_parser/core.py | 5 +- 2 files changed, 85 insertions(+), 26 deletions(-) diff --git a/database-lib/database_lib/main.py b/database-lib/database_lib/main.py index 5052f5a..7241a1b 100644 --- a/database-lib/database_lib/main.py +++ b/database-lib/database_lib/main.py @@ -24,7 +24,11 @@ class CacheData: self.data = data def json(self): - return json.loads(self.data) + try: + return json.loads(self.data) + except json.JSONDecodeError: + logger.warning(f"Failed to decode JSON data: {self.data[:100]}...") + return None def __repr__(self): return self.data @@ -35,6 +39,7 @@ class CacheData: def has_data(self): return self.data is not None and self.data != "" + class CacheResponse: __slots__ = ("key", "data") @@ -123,14 +128,18 @@ class SQLiteCacheBackend(AbstractCacheBackend): def random(self, size: int) -> list[CacheResponse]: self.ensure_connection() with self.connection: - self.cursor.execute("SELECT key, value FROM cache ORDER BY RANDOM() LIMIT ?", (size,)) + self.cursor.execute( + "SELECT key, value FROM cache ORDER BY RANDOM() LIMIT ?", (size,) + ) return [CacheResponse(key, value) for key, value in self.cursor] def enable_zstd(self): self.ensure_connection() with self.connection: try: - self.cursor.execute('SELECT zstd_enable_transparent(\'{"table": "cache", "column": "value", "compression_level": 9, "dict_chooser": "\'\'a\'\'"}\')') + self.cursor.execute( + 'SELECT zstd_enable_transparent(\'{"table": "cache", "column": "value", "compression_level": 9, "dict_chooser": "\'\'a\'\'"}\')' + ) except Exception as error: logger.error(f"Error enabling ZSTD compression: {error}") logger.exception(error) @@ -140,13 +149,17 @@ class SQLiteCacheBackend(AbstractCacheBackend): def init_db(self): self.ensure_connection() with self.connection: - self.cursor.execute("CREATE TABLE IF NOT EXISTS cache (key TEXT PRIMARY KEY, value TEXT)") + 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[CacheResponse, None]: self.ensure_connection() with self.connection: - cache = self.cursor.execute("SELECT value FROM cache WHERE key = :0", {"0": key}).fetchone() + 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(key, cache[0]) @@ -161,17 +174,24 @@ class SQLiteCacheBackend(AbstractCacheBackend): 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 a string or dict, not {type(value).__name__}") + raise ValueError( + f"value argument should be a string or dict, not {type(value).__name__}" + ) self.ensure_connection() with self.lock: with self.connection: - self.cursor.execute("INSERT OR REPLACE INTO cache VALUES (:0, :1)", {"0": key, "1": value}) + self.cursor.execute( + "INSERT OR REPLACE INTO cache VALUES (:0, :1)", + {"0": key, "1": value}, + ) def delete(self, key: str) -> None: self.ensure_connection() with self.connection: - result = self.cursor.execute("SELECT 1 FROM cache WHERE key = :0", {"0": key}).fetchone() + 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}") @@ -183,7 +203,9 @@ class SQLiteCacheBackend(AbstractCacheBackend): self.ensure_connection() with self.connection: # Fetch a random key-value pair just once - self.cursor.execute("SELECT key, value FROM cache ORDER BY RANDOM() LIMIT 1") + self.cursor.execute( + "SELECT key, value FROM cache ORDER BY RANDOM() LIMIT 1" + ) key, value = self.cursor.fetchone() # Use a generator to create batches of test data @@ -196,14 +218,20 @@ class SQLiteCacheBackend(AbstractCacheBackend): batch = list(islice(data_generator(), batch_size)) self.cursor.executemany("INSERT INTO cache VALUES (?, ?)", batch) self.connection.commit() - print(f"Inserted {min(offset + batch_size, num_rows)} / {num_rows} rows") + print( + f"Inserted {min(offset + batch_size, num_rows)} / {num_rows} rows" + ) def maintenance(self, time: Optional[int] = None, blocking_time: float = 0.5): connection = sqlite3.connect(self.database) cursor = connection.cursor() connection.enable_load_extension(True) # Enable loading of extensions - connection.execute("PRAGMA foreign_keys = ON;") # Need for working with foreign keys in db - connection.execute("PRAGMA journal_mode=WAL;") # Need to properly work with ZSTD compression + connection.execute( + "PRAGMA foreign_keys = ON;" + ) # Need for working with foreign keys in db + connection.execute( + "PRAGMA journal_mode=WAL;" + ) # Need to properly work with ZSTD compression connection.execute("PRAGMA auto_vacuum=full;") # Same as above thing if sqlite_zstd is not None: @@ -211,9 +239,13 @@ class SQLiteCacheBackend(AbstractCacheBackend): with connection: if time is not None: - cursor.execute("SELECT zstd_incremental_maintenance(?, ?);", (time, blocking_time)) + cursor.execute( + "SELECT zstd_incremental_maintenance(?, ?);", (time, blocking_time) + ) else: - cursor.execute("SELECT zstd_incremental_maintenance(null, ?);", (blocking_time,)) + cursor.execute( + "SELECT zstd_incremental_maintenance(null, ?);", (blocking_time,) + ) cursor.execute("VACUUM") cursor.execute("ANALYZE") @@ -283,7 +315,9 @@ class PostgreSQLCacheBackend(AbstractCacheBackend): def random(self, size: int) -> list[CacheResponse]: self.ensure_connection() with self.connection: - self.cursor.execute("SELECT key, value FROM cache ORDER BY RANDOM() LIMIT %s", (size,)) + self.cursor.execute( + "SELECT key, value FROM cache ORDER BY RANDOM() LIMIT %s", (size,) + ) return [CacheResponse(key, value) for key, value in self.cursor] def pull(self, key: str) -> Union[CacheResponse, None]: @@ -305,11 +339,16 @@ class PostgreSQLCacheBackend(AbstractCacheBackend): 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 a string or dict, not {type(value).__name__}") + raise ValueError( + f"value argument should be a string or dict, not {type(value).__name__}" + ) self.ensure_connection() with self.connection: - self.cursor.execute("INSERT INTO cache (key, value) VALUES (%s, %s) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", (key, value)) + self.cursor.execute( + "INSERT INTO cache (key, value) VALUES (%s, %s) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + (key, value), + ) def delete(self, key: str) -> None: self.ensure_connection() @@ -329,7 +368,9 @@ class PostgreSQLCacheBackend(AbstractCacheBackend): self.connection = None -def migrate_to_postgres(sqlite_db_path: str, pg_conn_string: str, chunk_size: int = 1000): +def migrate_to_postgres( + sqlite_db_path: str, pg_conn_string: str, chunk_size: int = 1000 +): logger.debug(f"Starting migration from SQLite ({sqlite_db_path}) to PostgreSQL") sqlite_db = SQLiteCacheBackend(sqlite_db_path, zstd_enabled=True) pg_db = PostgreSQLCacheBackend(pg_conn_string) @@ -345,20 +386,29 @@ def migrate_to_postgres(sqlite_db_path: str, pg_conn_string: str, chunk_size: in try: while processed_rows < total_rows: - chunk = sqlite_db.cursor.execute("SELECT key, value FROM cache LIMIT ? OFFSET ?", (chunk_size, processed_rows)).fetchall() + chunk = sqlite_db.cursor.execute( + "SELECT key, value FROM cache LIMIT ? OFFSET ?", + (chunk_size, processed_rows), + ).fetchall() if not chunk: logger.debug("No more rows to process") break logger.debug(f"Processing chunk of {len(chunk)} rows") - execute_batch(pg_db.cursor, "INSERT INTO cache (key, value) VALUES (%s, %s) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", chunk) + execute_batch( + pg_db.cursor, + "INSERT INTO cache (key, value) VALUES (%s, %s) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + chunk, + ) pg_db.connection.commit() processed_rows += len(chunk) elapsed_time = time.time() - start_time rows_per_second = processed_rows / elapsed_time - logger.info(f"Processed {processed_rows}/{total_rows} rows. Speed: {rows_per_second:.2f} rows/second") + logger.info( + f"Processed {processed_rows}/{total_rows} rows. Speed: {rows_per_second:.2f} rows/second" + ) except Exception as e: logger.error(f"An error occurred during migration: {e}") pg_db.connection.rollback() @@ -369,11 +419,19 @@ def migrate_to_postgres(sqlite_db_path: str, pg_conn_string: str, chunk_size: in total_time = time.time() - start_time logger.success(f"Data migration to PostgreSQL completed") - logger.info(f"Total time: {total_time:.2f} seconds. Average speed: {total_rows/total_time:.2f} rows/second") + logger.info( + f"Total time: {total_time:.2f} seconds. Average speed: {total_rows/total_time:.2f} rows/second" + ) -def execute_migrate_to_postgres_in_thread(sqlite_db_path: str, pg_conn_string: str, chunk_size: int = 1000): +def execute_migrate_to_postgres_in_thread( + sqlite_db_path: str, pg_conn_string: str, chunk_size: int = 1000 +): logger.info("Starting migration to PostgreSQL in thread") - migration_thread = threading.Thread(target=migrate_to_postgres, args=(sqlite_db_path, pg_conn_string, chunk_size), daemon=True) + migration_thread = threading.Thread( + target=migrate_to_postgres, + args=(sqlite_db_path, pg_conn_string, chunk_size), + daemon=True, + ) migration_thread.start() return migration_thread diff --git a/medium-parser/medium_parser/core.py b/medium-parser/medium_parser/core.py index 87af8e4..35fffba 100644 --- a/medium-parser/medium_parser/core.py +++ b/medium-parser/medium_parser/core.py @@ -101,10 +101,11 @@ class MediumParser: async def _get_from_cache(): logger.debug("Using cache backend") post_data = self.cache.pull(post_id) - logger.info(f"Found data in cache: {post_data[:10]}") if post_data and post_data.data.has_data(): logger.debug("post query was found on cache") - return post_data.json() + parsed_data = post_data.json() + if parsed_data: + return parsed_data logger.debug(f"No data found in cache by {post_id}") return None