From d9f574637c589cb7a5cb64feba3df91e6add2759 Mon Sep 17 00:00:00 2001
From: Eric Hare <ericrhare@gmail.com>
Date: Wed, 15 Nov 2023 14:40:33 -0800
Subject: [PATCH] Fix for upserting of documents

---
 astrapy/db.py            | 20 +++++++-------------
 tests/astrapy/test_db.py |  6 +++---
 2 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/astrapy/db.py b/astrapy/db.py
index 792b107c..4f3975c1 100644
--- a/astrapy/db.py
+++ b/astrapy/db.py
@@ -635,24 +635,18 @@ def upsert(self, document):
             str: The _id of the inserted or updated document.
         """
         # Attempt to insert the given document
-        result = self.insert_one(document)
-
-        # Check if we hit an error
-        if (
-            "errors" in result
-            and "errorCode" in result["errors"][0]
-            and result["errors"][0]["errorCode"] == "DOCUMENT_ALREADY_EXISTS"
-        ):
+        try:
+            self.insert_one(document)
+        except Exception as e:
+            logger.debug(e)
+
             # Now we attempt to update
-            result = self.find_one_and_replace(
+            self.find_one_and_replace(
                 filter={"_id": document["_id"]},
                 replacement=document,
             )
-            upserted_id = result["data"]["document"]["_id"]
-        else:
-            upserted_id = result["status"]["insertedIds"][0]
 
-        return upserted_id
+        return document["_id"]
 
 
 class AstraDB:
diff --git a/tests/astrapy/test_db.py b/tests/astrapy/test_db.py
index 919f48e7..c10402a3 100644
--- a/tests/astrapy/test_db.py
+++ b/tests/astrapy/test_db.py
@@ -600,7 +600,7 @@ def test_vector_find_projection(projection_collection):
 
 
 @pytest.mark.describe("upsert a document")
-def test_upsert_document(collection, cliff_uuid):
+def test_upsert_document(collection):
     new_uuid = str(uuid.uuid4())
 
     collection.upsert(
@@ -624,12 +624,12 @@ def test_upsert_document(collection, cliff_uuid):
 
     collection.upsert(
         {
-            "_id": cliff_uuid,
+            "_id": new_uuid,
             "addresses": {"work": {"city": "Everett", "state": "WA", "country": "USA"}},
         }
     )
 
-    document = collection.find_one(filter={"_id": cliff_uuid})
+    document = collection.find_one(filter={"_id": new_uuid})
 
     assert document is not None
     assert document["data"]["document"]["addresses"]["work"]["city"] == "Everett"