From 1b0399f28691e84b4deb60c81bf486f5fa0c51b3 Mon Sep 17 00:00:00 2001
From: Hongbo Wu <hongbo@omnivore.app>
Date: Wed, 23 Mar 2022 12:09:25 +0800
Subject: [PATCH] add assertion of migrated data in python script

---
 packages/api/elastic_migrate.py | 32 +++++++++++++++++++++++++++-----
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/packages/api/elastic_migrate.py b/packages/api/elastic_migrate.py
index 3b625e3ac..ebabb2a63 100644
--- a/packages/api/elastic_migrate.py
+++ b/packages/api/elastic_migrate.py
@@ -64,6 +64,26 @@ UPDATE omnivore.highlight
 '''
 
 
+def assertData(conn, client):
+    # get all users from postgres
+    try:
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        cursor.execute('''SELECT id FROM omnivore.user''')
+        result = cursor.fetchall()
+        for row in result:
+            userId = row['id']
+            cursor.execute(
+                f'''SELECT COUNT(id) FROM omnivore.links WHERE user_id = "{userId}"''')
+            countInPostgres = cursor.fetchone()
+            countInElastic = client.count(
+                index='pages', body={'query': {'term': {'userId': userId}}})['count']
+            assert countInPostgres == countInElastic
+        cursor.close()
+    except Exception as err:
+        print('Assert data ERROR:', err)
+        exit(1)
+
+
 def create_index(client):
     print('Creating index')
     try:
@@ -183,15 +203,17 @@ conn = psycopg2.connect(
     password={PG_PASSWORD}')
 print('Postgres connection:', conn.info)
 
-create_index(client)
+# create_index(client)
 
 # ingest data from postgres to es and json file (for debugging)
-ingest_data_to_elastic(conn, QUERY, DATA_FILE)
+# ingest_data_to_elastic(conn, QUERY, DATA_FILE)
 
 # update existing tables
-update_postgres_data(conn, UPDATE_ARTICLE_SAVING_REQUEST_SQL,
-                     'article_saving_request')
-update_postgres_data(conn, UPDATE_HIGHLIGHT_SQL, 'highlight')
+# update_postgres_data(conn, UPDATE_ARTICLE_SAVING_REQUEST_SQL,
+#                      'article_saving_request')
+# update_postgres_data(conn, UPDATE_HIGHLIGHT_SQL, 'highlight')
+
+assertData(conn, client)
 
 client.close()
 conn.close()