db cache: fixing broken json parsing

This commit is contained in:
ZhymabekRoman 2024-09-21 16:18:22 +05:00
parent 8b109ff4ec
commit 14bbf1c1c5
2 changed files with 85 additions and 26 deletions

View file

@ -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

View file

@ -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