From 101825df4614a4630e9e6d38e3f485a07a65a1ee Mon Sep 17 00:00:00 2001 From: Hongbo Wu Date: Tue, 19 Sep 2023 17:54:25 +0800 Subject: [PATCH] add assertion of migrated data in python script --- .../migrate_from_elastic.py | 122 ++++++++++++++---- 1 file changed, 94 insertions(+), 28 deletions(-) diff --git a/packages/db/elastic_migrations/migrate_from_elastic.py b/packages/db/elastic_migrations/migrate_from_elastic.py index 8e0a93302..5c2437546 100644 --- a/packages/db/elastic_migrations/migrate_from_elastic.py +++ b/packages/db/elastic_migrations/migrate_from_elastic.py @@ -4,11 +4,11 @@ import hashlib import os import uuid from datetime import datetime + import asyncpg from elasticsearch import AsyncElasticsearch from elasticsearch.helpers import async_scan - PG_HOST = os.getenv('PG_HOST', 'localhost') PG_PORT = os.getenv('PG_PORT', 5432) PG_USER = os.getenv('PG_USER', 'hongbowu') @@ -25,6 +25,70 @@ CUT_OFF_DATE = os.getenv('CUT_OFF_DATE', '2000-01-01') DATE_FORMAT = '%Y-%m-%dT%H:%M:%S.%fZ' +async def assert_data(db_conn, es_client, users, uploaded_files): + # get all users from postgres + try: + success = 0 + failure = 0 + for user in users: + user_id = user['id'] + number_of_docs_in_postgres = await db_conn.fetchval( + f'SELECT COUNT(*) FROM omnivore.library_item WHERE user_id = \'{user_id}\'') + + query = { + 'size': 0, + 'query': { + 'bool': { + 'must': [ + { + 'term': { + 'userId': user_id + } + } + ], + 'should': [ + { + 'bool': { + 'must_not': { + 'exists': { + 'field': 'uploadFileId' + } + } + } + }, + { + 'terms': { + 'uploadFileId': [file['id'] for file in uploaded_files] + } + } + ], + 'minimum_should_match': 1 + } + }, + 'aggs': { + 'unique_urls': { + 'cardinality': { + 'field': 'url' + } + } + } + } + # Get the count of unique URLs for the specified user + result = await es_client.search(index=ES_INDEX, body=query) + number_of_docs_in_elastic = result['aggregations']['unique_urls']['value'] + + if number_of_docs_in_postgres == number_of_docs_in_elastic: + success += 1 + print(f'User {user_id} OK') + else: + failure += 1 + print( + f'User {user_id} ERROR: postgres: {number_of_docs_in_postgres}, elastic: {number_of_docs_in_elastic}') + print(f'Asserted data, success: {success}, failure: {failure}') + except Exception as err: + print('Assert data ERROR:', err) + + def get_uuid(val): try: uuid.UUID(val, version=4) @@ -46,7 +110,7 @@ def convert_string_to_datetime(val): async def insert_library_items(db_conn, library_items): - insert_query = """ + insert_query = ''' INSERT INTO omnivore.library_item ( id, user_id, title, author, description, readable_content, original_url, upload_file_id, item_type, slug, reading_progress_top_percent, reading_progress_bottom_percent, reading_progress_highest_read_anchor, created_at, saved_at, archived_at, site_name, subscription, state, updated_at, published_at, item_language, read_at, word_count, site_icon, thumbnail, content_reader, original_content ) @@ -79,14 +143,14 @@ async def insert_library_items(db_conn, library_items): thumbnail = EXCLUDED.thumbnail, content_reader = EXCLUDED.content_reader, original_content = EXCLUDED.original_content - """ + ''' print('Inserting library items into postgres') await insert_into_postgres(insert_query, db_conn, library_items) print(f'Inserted {len(library_items)} library items') async def insert_highlights(db_conn, highlights): - insert_query = """ + insert_query = ''' INSERT INTO omnivore.highlight ( id, user_id, quote, prefix, suffix, patch, annotation, created_at, updated_at, shared_at, short_id, library_item_id, highlight_position_percent, highlight_position_anchor_index, highlight_type, color, html ) @@ -114,14 +178,14 @@ async def insert_highlights(db_conn, highlights): highlight_type = EXCLUDED.highlight_type, color = EXCLUDED.color, html = EXCLUDED.html - """ + ''' print('Inserting highlights into postgres') await insert_into_postgres(insert_query, db_conn, highlights) print(f'Inserted {len(highlights)} highlights') async def insert_labels(db_conn, labels): - insert_query = """ + insert_query = ''' INSERT INTO omnivore.entity_labels ( label_id, library_item_id, highlight_id ) @@ -131,14 +195,14 @@ async def insert_labels(db_conn, labels): LEFT JOIN omnivore.highlight h ON h.id = $3 WHERE l.id = $1 AND ($2 IS NULL OR li.id = $2) AND ($3 IS NULL OR h.id = $3) ON CONFLICT (label_id, library_item_id, highlight_id) DO NOTHING - """ + ''' print('Inserting labels into postgres') await insert_into_postgres(insert_query, db_conn, labels) print(f'Inserted {len(labels)} labels') async def insert_recommendations(db_conn, recommendations): - insert_query = """ + insert_query = ''' INSERT INTO omnivore.recommendation ( library_item_id, recommender_id, group_id, note, created_at ) @@ -150,7 +214,7 @@ async def insert_recommendations(db_conn, recommendations): ON CONFLICT (library_item_id, recommender_id, group_id) DO UPDATE SET note = EXCLUDED.note, created_at = EXCLUDED.created_at - """ + ''' print('Inserting recommendations into postgres') await insert_into_postgres(insert_query, db_conn, recommendations) print(f'Inserted {len(recommendations)} recommendations') @@ -196,45 +260,45 @@ async def main(): recommendations = [] query = { - "query": { - "bool": { - "must": [ + 'query': { + 'bool': { + 'must': [ { - "range": { - "updatedAt": { - "gte": CUT_OFF_DATE + 'range': { + 'updatedAt': { + 'gte': CUT_OFF_DATE } } }, { - "terms": { - "userId": [user['id'] for user in users] + 'terms': { + 'userId': [user['id'] for user in users] }, }, ], - "should": [ + 'should': [ { - "bool": { - "must_not": { - "exists": { - "field": "uploadFileId" + 'bool': { + 'must_not': { + 'exists': { + 'field': 'uploadFileId' } } }, }, { - "terms": { - "uploadFileId": [file['id'] for file in uploaded_files] + 'terms': { + 'uploadFileId': [file['id'] for file in uploaded_files] } } ], - "minimum_should_match": 1, + 'minimum_should_match': 1, } }, - "sort": [ + 'sort': [ { - "updatedAt": { - "order": "asc" + 'updatedAt': { + 'order': 'asc' } } ] @@ -372,6 +436,8 @@ async def main(): await insert_recommendations(db_conn, recommendations) print('Migration complete', migrated_at) + + await assert_data(db_conn, es_client, users, uploaded_files) except Exception as err: print('Migration error', err) finally: