implement get_document_id/get_document_ids

Author: Reggie Burnett

lib/mysqlx/connection.py

@@ -30,7 +30,7 @@
 from .crud import Schema
 from .protocol import Protocol, MessageReaderWriter
 from .result import Result, RowResult, DocResult
-from .statement import SqlStatement
+from .statement import SqlStatement, AddStatement
 
 _DROP_DATABASE_QUERY = "DROP DATABASE IF EXISTS `{0}`"
 _CREATE_DATABASE_QUERY = "CREATE DATABASE IF NOT EXISTS `{0}`"
@@ -95,7 +95,10 @@ def send_sql(self, sql, *args):
 
     def send_insert(self, statement):
         self.protocol.send_insert(statement)
-        return Result(self)
+        ids = None
+        if isinstance(statement, AddStatement):
+            ids = statement._ids
+        return Result(self, ids)
 
     def find(self, statement):
         self.protocol.send_find(statement)

lib/mysqlx/dbdoc.py

@@ -54,6 +54,7 @@ def keys(self):
     def ensure_id(self):
         if "_id" not in self.__dict__:
             self.__dict__["_id"] = str(uuid.uuid4())
+        return self.__dict__["_id"]
 
     def __str__(self):
         return json.dumps(self.__dict__)

lib/mysqlx/result.py

@@ -533,8 +533,9 @@ def get_warnings_count(self):
 
 
 class Result(BaseResult):
-    def __init__(self, connection):
+    def __init__(self, connection, ids=None):
         super(Result, self).__init__(connection)
+        self._ids = ids
         self._protocol.close_result(self)
 
     @property
@@ -544,6 +545,14 @@ def get_affected_items_count(self):
     def get_autoincrement_value(self):
         return self._generated_id
 
+    def get_document_id(self):
+        if self._ids == None: return None
+        if len(self._ids) == 0: return None
+        return self._ids[0]
+
+    def get_document_ids(self):
+        return self._ids
+
 
 class BufferingResult(BaseResult):
     def __init__(self, connection):

lib/mysqlx/statement.py

@@ -204,6 +204,7 @@ class AddStatement(Statement):
     def __init__(self, collection):
         super(AddStatement, self).__init__(target=collection)
         self._values = []
+        self._ids = []
 
     """Adds a list of documents into a collection.
 
@@ -228,7 +229,8 @@ def execute(self):
             mysqlx.Result: Result object.
         """
         for doc in self._values:
-            doc.ensure_id()
+            self._ids.append(doc.ensure_id())
+
         return self._connection.send_insert(self)
 
 

tests/test_mysqlx_crud.py

@@ -181,6 +181,17 @@ def test_add(self):
         self.assertEqual(result.get_affected_items_count, 2)
         self.assertEqual(5, collection.count())
 
+    def test_get_document_ids(self):
+        collection_name = "collection_test"
+        collection = self.schema.create_collection(collection_name)
+        result = collection.add({"name": "Fred", "age": 21}).execute()
+        self.assertTrue(result.get_document_id() != None)
+
+        result = collection.add(
+            {"name": "Fred", "age": 21},
+            {"name": "Barney", "age":45}).execute()
+        self.assertEqual(2, len(result.get_document_ids()))
+
     def test_remove(self):
         collection_name = "collection_test"
         collection = self.schema.create_collection(collection_name)