Merge pull request #683 from praekeltfoundation/update-migration-script

Update migration script

Author: erikh360

scripts/migrate_to_turn/fetch_rapidpro_contacts.py

@@ -20,22 +20,25 @@
     to_lowercase,
 )
 
-RAPIDPRO_URL = "https://rapidpro.qa.momconnect.co.za"
-
-START_DATE = "2026-02-24 01:13:06"
-END_DATE = "2026-02-26 19:13:06"
-LIMIT = 1000
+env = "qa"  # qa or prd
+RAPIDPRO_URL = f"https://rapidpro.{env}.momconnect.co.za"
+
+START_DATE = "2022-02-10 00:0:00"
+# END_DATE = "2026-02-26 19:13:06"
+# START_DATE = "2026-03-04 00:00:00"
+END_DATE = "2026-03-05 00:00:00"
+LIMIT = 1000000
 INCLUDE_OPTED_OUT = False
 # We want to import beta testing users as opted out and give them the chance to opt in.
 IMPORT_AS_OPTED_OUT = False
 # This is to identify invited users and schedule the invite message.
-MIGRATION_KEY = "beta_testing_batch_1"
+MIGRATION_KEY = "batch_1"
 
-MSISDN_FILTER = (
-    os.environ.get("MSISDN_FILTER", "").split(",")
-    if os.environ.get("MSISDN_FILTER")
-    else []
-)
+MSISDN_FILTER = [
+    msisdn.strip()
+    for msisdn in os.environ.get("MSISDN_FILTER", "").split(",")
+    if msisdn.strip()
+]
 
 FIELD_MAPPING = {
     "edd": {
@@ -108,6 +111,15 @@
 }
 
 
+def get_readable_timestamp(dt=None):
+    dt = dt or datetime.now().astimezone()
+    return dt.strftime("%Y-%m-%d %H:%M:%S %Z")
+
+
+def log(message):
+    print(f"[{get_readable_timestamp()}] {message}")
+
+
 def get_field_data(contact):
     data = {}
     for rapidpro_field, turn_details in FIELD_MAPPING.items():
@@ -118,6 +130,11 @@ def get_field_data(contact):
         else:
             data[turn_field] = contact.fields.get(rapidpro_field)
 
+        if rapidpro_field == "language":
+            language_value = data[turn_field]
+            if language_value is None or str(language_value).strip() == "":
+                data[turn_field] = "eng"
+
         if "process" in turn_details:
             data[turn_field] = turn_details["process"](data[turn_field])
 
@@ -138,53 +155,108 @@ def is_opted_out(contact):
     return opted_out.upper() == "TRUE"
 
 
-def get_rapidpro_contacts(client, start_date=None, end_date=None):
-    print(f"> Getting rapidpro contacts from {start_date} to {end_date}")
-    contacts = []
-    oldest_date = end_date.replace(tzinfo=pytz.utc)
-    for contact_batch in client.get_contacts(
-        before=end_date, after=start_date
-    ).iterfetches(retry_on_rate_exceed=True):
-        for contact in contact_batch:
-            wa_id = get_wa_id(contact)
+def get_contact_urn(msisdn):
+    return msisdn if msisdn.startswith("whatsapp:") else f"whatsapp:{msisdn}"
+
+
+def process_contact(
+    contact, contacts, oldest_date, seen_wa_ids=None, apply_msisdn_filter=False
+):
+    wa_id = get_wa_id(contact)
+
+    if is_opted_out(contact) and not INCLUDE_OPTED_OUT:
+        return oldest_date, False
+
+    if not wa_id:
+        return oldest_date, False
+
+    if apply_msisdn_filter and MSISDN_FILTER and wa_id not in MSISDN_FILTER:
+        return oldest_date, False
 
-            if is_opted_out(contact) and not INCLUDE_OPTED_OUT:
-                continue
+    if seen_wa_ids is not None and wa_id in seen_wa_ids:
+        return oldest_date, False
 
-            # TODO: filtering for testing only, remove later()
-            if wa_id and MSISDN_FILTER and wa_id not in MSISDN_FILTER:
-                continue
+    data = get_field_data(contact)
+    data["urn"] = wa_id
+    contacts.append(data)
 
-            if wa_id:
-                data = get_field_data(contact)
-                data["urn"] = wa_id
-                contacts.append(data)
+    if seen_wa_ids is not None:
+        seen_wa_ids.add(wa_id)
 
-                modified_on = contact.modified_on.astimezone(pytz.utc)
-                if modified_on < oldest_date:
-                    oldest_date = modified_on
+    modified_on = contact.modified_on.astimezone(pytz.utc)
+    if oldest_date is None or modified_on < oldest_date:
+        oldest_date = modified_on
 
-                if len(contacts) >= LIMIT:
+    return oldest_date, True
+
+
+def get_rapidpro_contacts_by_msisdn(client):
+    log(f"> Getting rapidpro contacts by MSISDN filter ({len(MSISDN_FILTER)} values)")
+    contacts = []
+    oldest_date = None
+    seen_wa_ids = set()
+    batch_number = 0
+
+    for msisdn in MSISDN_FILTER:
+        urn = get_contact_urn(msisdn)
+        for contact_batch in client.get_contacts(urn=urn).iterfetches(
+            retry_on_rate_exceed=True
+        ):
+            batch_number += 1
+            log(f"Processing contact_batch #{batch_number} for {urn}")
+            for contact in contact_batch:
+                oldest_date, added = process_contact(
+                    contact, contacts, oldest_date, seen_wa_ids=seen_wa_ids
+                )
+
+                if added and len(contacts) >= LIMIT:
                     return contacts, oldest_date
 
     return contacts, oldest_date
 
 
-def fetch_rapidpro_contacts(client):
-    start_date = datetime.strptime(START_DATE, "%Y-%m-%d %H:%M:%S")
-    end_date = datetime.strptime(END_DATE, "%Y-%m-%d %H:%M:%S")
+def get_rapidpro_contacts(client, start_date=None, end_date=None):
+    log(f"> Getting rapidpro contacts from {start_date} to {end_date}")
+    contacts = []
+    oldest_date = end_date.replace(tzinfo=pytz.utc)
+    for batch_number, contact_batch in enumerate(
+        client.get_contacts(before=end_date, after=start_date).iterfetches(
+            retry_on_rate_exceed=True
+        ),
+        start=1,
+    ):
+        log(f"Processing contact_batch #{batch_number}")
+        for contact in contact_batch:
+            oldest_date, added = process_contact(
+                contact, contacts, oldest_date, apply_msisdn_filter=True
+            )
 
-    contacts, oldest_date = get_rapidpro_contacts(client, start_date, end_date)
+            if added and len(contacts) >= LIMIT:
+                return contacts, oldest_date
 
-    print(f"Found: {len(contacts)}")
-    print(f"Oldest modified on date: {oldest_date}")
+    return contacts, oldest_date
 
-    if contacts:
+
+def fetch_rapidpro_contacts(client):
+    run_start = datetime.now().astimezone()
+    log("Starting RapidPro contact fetch")
+
+    if MSISDN_FILTER:
+        contacts, oldest_date = get_rapidpro_contacts_by_msisdn(client)
+        filename = f"contacts-{env}-msisdn-filter.csv"
+    else:
+        start_date = datetime.strptime(START_DATE, "%Y-%m-%d %H:%M:%S")
+        end_date = datetime.strptime(END_DATE, "%Y-%m-%d %H:%M:%S")
+        contacts, oldest_date = get_rapidpro_contacts(client, start_date, end_date)
         start = START_DATE.split(" ")[0]
         end = END_DATE.split(" ")[0]
-        filename = f"contacts-{start}-{end}.csv"
+        filename = f"contacts-{env}-{start}-{end}.csv"
 
-        print(f"File: {filename}")
+    log(f"Found: {len(contacts)}")
+    log(f"Oldest modified on date: {oldest_date}")
+
+    if contacts:
+        log(f"File: {filename}")
 
         keys = contacts[0].keys()
 
@@ -193,6 +265,12 @@ def fetch_rapidpro_contacts(client):
             dict_writer.writeheader()
             dict_writer.writerows(contacts)
 
+    run_end = datetime.now().astimezone()
+    log("Finished RapidPro contact fetch")
+    log(f"Run started: {get_readable_timestamp(run_start)}")
+    log(f"Run ended: {get_readable_timestamp(run_end)}")
+    log(f"Total duration: {run_end - run_start}")
+
 
 if __name__ == "__main__":
     client = TembaClient(RAPIDPRO_URL, os.environ["RAPIDPRO_TOKEN"])

scripts/migrate_to_turn/tests/test_fetch_rapidpro_contacts.py

@@ -34,8 +34,19 @@ class FakeClient:
     def __init__(self, batches):
         self._batches = batches
 
-    def get_contacts(self, before=None, after=None):
-        return FakeContacts(self._batches)
+    def get_contacts(self, before=None, after=None, urn=None):
+        if urn is None:
+            return FakeContacts(self._batches)
+
+        filtered_batches = []
+        for batch in self._batches:
+            filtered_batch = [
+                contact for contact in batch if urn in getattr(contact, "urns", [])
+            ]
+            if filtered_batch:
+                filtered_batches.append(filtered_batch)
+
+        return FakeContacts(filtered_batches)
 
 
 class FetchRapidproContactsTests(TestCase):
@@ -174,7 +185,7 @@ def test_fetch_rapidpro_contacts_writes_csv(self):
             ):
                 fetch_rapidpro_contacts.fetch_rapidpro_contacts(client)
 
-            output_path = Path(tmp_dir) / "contacts-2025-01-01-2025-01-31.csv"
+            output_path = Path(tmp_dir) / "contacts-qa-2025-01-01-2025-01-31.csv"
             with output_path.open(newline="") as csv_file:
                 reader = csv.DictReader(csv_file)
                 rows = list(reader)
@@ -268,7 +279,7 @@ def test_fetch_rapidpro_contacts_skips_opted_out(self):
             ):
                 fetch_rapidpro_contacts.fetch_rapidpro_contacts(client)
 
-            output_path = Path(tmp_dir) / "contacts-2025-01-01-2025-01-31.csv"
+            output_path = Path(tmp_dir) / "contacts-qa-2025-01-01-2025-01-31.csv"
             self.assertFalse(output_path.exists())
 
     def test_fetch_rapidpro_contacts_msisdn_filter(self):
@@ -310,7 +321,7 @@ def test_fetch_rapidpro_contacts_msisdn_filter(self):
             ):
                 fetch_rapidpro_contacts.fetch_rapidpro_contacts(client)
 
-            output_path = Path(tmp_dir) / "contacts-2025-01-01-2025-01-31.csv"
+            output_path = Path(tmp_dir) / "contacts-qa-msisdn-filter.csv"
             with output_path.open(newline="") as csv_file:
                 reader = csv.DictReader(csv_file)
                 rows = list(reader)