allow replace (#39)

dolthub · Jun 29, 2022 · fd79aa2 · fd79aa2
1 parent 6d916da
commit fd79aa2
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/doltcli/utils.py b/doltcli/utils.py
@@ -245,12 +245,12 @@ def _get_import_mode_and_flags(
     import_modes = IMPORT_MODES_TO_FLAGS.keys()
     if import_mode and import_mode not in import_modes:
         raise ValueError(f"update_mode must be one of: {import_modes}")
-    else:
+    elif not import_mode:
         if table in [table.name for table in dolt.ls()]:
             logger.info(f'No import mode specified, table exists, using "{UPDATE}"')
             import_mode = UPDATE
         else:
-            logger.info(f'No import mode specified, table exists, using "{CREATE}"')
+            logger.info(f'No import mode specified, table does not exist, using "{CREATE}"')
             import_mode = CREATE
 
     return import_mode

diff --git a/tests/test_write.py b/tests/test_write.py
@@ -35,6 +35,36 @@ def test_write_rows(init_empty_test_repo):
     compare_rows_helper(TEST_ROWS, actual)
 
 
+def test_update_rows(init_empty_test_repo):
+    dolt = init_empty_test_repo
+    write_rows(dolt, "characters", TEST_ROWS, CREATE, ["id"])
+
+    new_row = {"name": "dick butkus", "adjective": "buffoon", "id": "3", "date_of_death": ""}
+
+    write_rows(dolt, "characters", [new_row], "update", ["id"])
+    actual = read_rows(dolt, "characters")
+    exp = [
+        {"name": "Anna", "adjective": "tragic", "id": "1", "date_of_death": "1877-01-01"},
+        {"name": "Vronksy", "adjective": "honorable", "id": "2", "date_of_death": ""},
+        {"name": "dick butkus", "adjective": "buffoon", "id": "3", "date_of_death": ""},
+    ]
+    compare_rows_helper(exp, actual)
+
+
+def test_replace_rows(init_empty_test_repo):
+    dolt = init_empty_test_repo
+    write_rows(dolt, "characters", TEST_ROWS, CREATE, ["id"])
+
+    new_row = {"name": "dick butkus", "adjective": "buffoon", "id": "3", "date_of_death": ""}
+
+    write_rows(dolt, "characters", [new_row], "replace", ["id"])
+    actual = read_rows(dolt, "characters")
+    exp = [
+        {"name": "dick butkus", "adjective": "buffoon", "id": "3", "date_of_death": ""},
+    ]
+    compare_rows_helper(exp, actual)
+
+
 def test_write_columns(init_empty_test_repo):
     dolt = init_empty_test_repo
     write_columns(dolt, "characters", TEST_COLUMNS, CREATE, ["id"])