Merge pull request #1156 from facebookresearch/data-porter-feature-un…

…ittests Added unittests for Data Porter command options
facebookresearch · May 13, 2024 · 4c0a0ba · 4c0a0ba
2 parents df76170 + 43aa665
commit 4c0a0ba
Show file tree

Hide file tree

Showing 5 changed files with 955 additions and 45 deletions.
diff --git a/mephisto/tools/db_data_porter/constants.py b/mephisto/tools/db_data_porter/constants.py
@@ -23,6 +23,7 @@
 ]
 DATASTORE_EXPORT_METHOD_NAME = "get_export_data"
 DEFAULT_CONFLICT_RESOLVER = "DefaultMergeConflictResolver"
+EXAMPLE_CONFLICT_RESOLVER = "ExampleMergeConflictResolver"
 IMPORTED_DATA_TABLE_NAME = "imported_data"
 MIGRATIONS_TABLE_NAME = "migrations"
 TASK_RUNS_TABLE_NAME = "task_runs"

diff --git a/mephisto/tools/db_data_porter/db_data_porter.py b/mephisto/tools/db_data_porter/db_data_porter.py
@@ -47,7 +47,7 @@ class DBDataPorter:
     This class contains the main logic of commands `mephisto db ...`.
     """
 
-    def __init__(self, db=None):
+    def __init__(self, db: "MephistoDB" = None):
         # Load Mephisto DB and providers' datastores
         if db is None:
             db = LocalMephistoDB()
@@ -324,7 +324,7 @@ def import_dump(
 
         if not os.path.exists(dump_archive_file_name_or_path):
             error_message = (
-                f"Could not find dump file '{dump_archive_file_name_or_path}'. "
+                f"Could not find dump file {dump_archive_file_name_or_path}. "
                 f"Please specify full path to existing file or "
                 f"only filename if located in <MEPHISTO_REPO>/{EXPORT_OUTPUT_DIR}."
             )
@@ -342,7 +342,7 @@ def import_dump(
                     dump_file_data: dict = json.loads(f.read())
                 except Exception as e:
                     error_message = (
-                        f"Could not read JSON from dump file '{dump_archive_file_name_or_path}'. "
+                        f"Could not read JSON from dump file {dump_archive_file_name_or_path}. "
                         f"Please, check if file '{json_dump_file_name}' in it "
                         f"has the correct format. Reason: {str(e)}"
                     )
@@ -377,7 +377,7 @@ def import_dump(
         dump_timestamp = self._make_export_timestamp()
         backup_path = backups.make_backup_file_path_by_timestamp(backup_dir, dump_timestamp)
         backups.make_full_data_dir_backup(backup_path)
-        logger.info(f"Backup was created successfully! File: '{backup_path}'")
+        logger.info(f"Backup was created successfully! File: {backup_path}")
 
         # 7. Write dump data into local DBs
         logger.info(f"Started importing from dump file {dump_archive_file_name_or_path} ...")

diff --git a/mephisto/tools/db_data_porter/randomize_ids.py b/mephisto/tools/db_data_porter/randomize_ids.py
@@ -168,7 +168,11 @@ def randomize_ids(
     # Providers' DBs
     provider_types = [i["provider_type"] for i in mephisto_dump["requesters"]]
     for provider_type in provider_types:
-        provider_dump = full_dump[provider_type]
+        provider_dump = full_dump.get(provider_type)
+
+        if not provider_dump:
+            continue
+
         randomized_ids_for_provider = _randomize_ids_for_provider(
             provider_type,
             provider_dump,

diff --git a/mephisto/utils/db.py b/mephisto/utils/db.py
@@ -572,7 +572,7 @@ def select_fk_mappings_for_tables(db: "MephistoDB", table_names: List[str]) -> d
     return tables_fks
 
 
-def insert_new_row_in_table(db: "MephistoDB", table_name: str, row: dict):
+def insert_new_row_in_table(db: "MephistoDB", table_name: str, row: dict) -> str:
     with db.table_access_condition, db.get_connection() as conn:
         c = conn.cursor()
 
@@ -590,6 +590,9 @@ def insert_new_row_in_table(db: "MephistoDB", table_name: str, row: dict):
             values,
         )
 
+        _id = str(c.lastrowid)
+        return _id
+
 
 def update_row_in_table(
     db: "MephistoDB",