WL10975: Add Single document operations

These are commands that operate at a single document level, unlike the other
CRUD commands that operate on all documents that match a filter.

This feature adds 4 methods to `Collection`:
1. `Collection.replace_one` replaces a document identified with a document_id
   with the provided document.
2. `Collection.add_or_replace_one` does the same as replace_one, but this adds
   a new document if the document_id doesn't exist.
3. `Collection.get_one` returns a document matching a document_id.
4. `Collection.remove_one` removes a document matching a document_id.

This patch also adds the upsert(val=True) method chainable to Table.insert() and
Collection.add().

By default upsert is set to False. Setting this to True replaces the row matched
with the provided key with the values provided.

Tests have been added for regression.

Author: amitab

lib/mysqlx/__init__.py

@@ -45,7 +45,7 @@
                         CreateCollectionIndexStatement, CreateTableStatement,
                         CreateViewStatement, AlterViewStatement, ColumnDef,
                         GeneratedColumnDef, ForeignKeyDef, Expr,
-                        ReadStatement)
+                        ReadStatement, WriteStatement)
 
 _SPLIT = re.compile(r',(?![^\(\)]*\))')
 _PRIORITY = re.compile(r'^\(address=(.+),priority=(\d+)\)$', re.VERBOSE)

lib/mysqlx/crud.py

@@ -23,6 +23,7 @@
 
 """Implementation of the CRUD database objects."""
 
+from .dbdoc import DbDoc
 from .errors import ProgrammingError
 from .statement import (FindStatement, AddStatement, RemoveStatement,
                         ModifyStatement, SelectStatement, InsertStatement,
@@ -409,6 +410,46 @@ def drop_index(self, index_name):
         self._connection.execute_nonquery("xplugin", "drop_collection_index",
             False, self._schema.name, self._name, index_name)
 
+    def replace_one(self, doc_id, doc):
+        """Replaces the Document matching the document ID with a
+           new document provided.
+
+        Args:
+            doc_id (str): Document ID
+            doc (DbDoc/dict): New Document
+        """
+        return self.modify("_id = :id").set("$", doc) \
+                   .bind("id", doc_id).execute()
+
+    def add_or_replace_one(self, doc_id, doc):
+        """Upserts the Document matching the document ID with a
+           new document provided.
+
+        Args:
+            doc_id (str): Document ID
+            doc (DbDoc/dict): New Document
+        """
+        if not isinstance(doc, DbDoc):
+            doc = DbDoc(doc)
+        doc.ensure_id(doc_id)
+        return self.add(doc).upsert(True).execute()
+
+    def get_one(self, doc_id):
+        """Returns a Document matching the Document ID.
+
+        Args:
+            doc_id (str): Document ID
+        """
+        return self.find("_id = :id").bind("id", doc_id).execute().fetch_one()
+
+    def remove_one(self, doc_id):
+        """Removes a Document matching the Document ID.
+
+        Args:
+            doc_id (str): Document ID
+        """
+        return self.remove("_id = :id").bind("id", doc_id).execute()
+
 
 class Table(DatabaseObject):
     """Represents a database table on a schema.

lib/mysqlx/dbdoc.py

@@ -27,6 +27,7 @@
 import uuid
 
 from .compat import STRING_TYPES
+from .errors import ProgrammingError
 
 
 class DbDoc(object):
@@ -47,14 +48,21 @@ def __init__(self, value):
         else:
             raise ValueError("Unable to handle type: {0}".format(type(value)))
 
+    def __setitem__(self, index, value):
+        if index == "_id":
+            raise ProgrammingError("Cannot modify _id")
+        self.__dict__[index] = value
+
     def __getitem__(self, index):
         return self.__dict__[index]
 
     def keys(self):
         return self.__dict__.keys()
 
-    def ensure_id(self):
-        if "_id" not in self.__dict__:
+    def ensure_id(self, doc_id=None):
+        if doc_id:
+            self.__dict__["_id"] = doc_id
+        elif "_id" not in self.__dict__:
             uuid1 = str(uuid.uuid1()).upper().split("-")
             uuid1.reverse()
             self.__dict__["_id"] = "".join(uuid1)

lib/mysqlx/expr.py

@@ -21,8 +21,9 @@
 # along with this program; if not, write to the Free Software
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
 
-from .compat import STRING_TYPES, PY3
+from .compat import STRING_TYPES, BYTE_TYPES, UNICODE_TYPES, PY3
 from .helpers import get_item_or_attr
+from .dbdoc import DbDoc
 from .protobuf import Message, mysqlxpb_enum
 
 
@@ -269,6 +270,51 @@ def __str__(self):
 
 # static protobuf helper functions
 
+def build_expr(value):
+    msg = Message("Mysqlx.Expr.Expr")
+    if isinstance(value, (dict, DbDoc)):
+        msg["type"] = mysqlxpb_enum("Mysqlx.Expr.Expr.Type.OBJECT")
+        msg["object"] = build_object(value).get_message()
+    elif isinstance(value, (list, tuple)):
+        msg["type"] = mysqlxpb_enum("Mysqlx.Expr.Expr.Type.ARRAY")
+        msg["array"] = build_array(value).get_message()
+    else:
+        msg["type"] = mysqlxpb_enum("Mysqlx.Expr.Expr.Type.LITERAL")
+        msg["literal"] = build_scalar(value).get_message()
+    return msg
+
+def build_scalar(value):
+    if isinstance(value, STRING_TYPES):
+        return build_string_scalar(value)
+    elif isinstance(value, BYTE_TYPES):
+        return build_bytes_scalar(value)
+    elif isinstance(value, bool):
+        return build_bool_scalar(value)
+    elif isinstance(value, int):
+        return build_int_scalar(value)
+    elif isinstance(value, float):
+        return build_double_scalar(value)
+    elif value is None:
+        return build_null_scalar()
+    raise ValueError("Unsupported data type: {0}.".format(type(value)))
+
+def build_object(obj):
+    if isinstance(obj, DbDoc):
+        return build_object(obj.__dict__)
+
+    msg = Message("Mysqlx.Expr.Object")
+    for key, value in obj.items():
+        pair = Message("Mysqlx.Expr.Object.ObjectField")
+        pair["key"] = key.encode() if isinstance(key, UNICODE_TYPES) else key
+        pair["value"] = build_expr(value).get_message()
+        msg["fld"].extend([pair.get_message()])
+    return msg
+
+def build_array(array):
+    msg = Message("Mysqlx.Expr.Array")
+    msg["value"].extend([build_expr(value).get_message() for value in array])
+    return msg
+
 def build_null_scalar():
     msg = Message("Mysqlx.Datatypes.Scalar")
     msg["type"] = mysqlxpb_enum("Mysqlx.Datatypes.Scalar.Type.V_NULL")

lib/mysqlx/protocol.py

@@ -28,8 +28,7 @@
 from .compat import STRING_TYPES, INT_TYPES
 from .dbdoc import DbDoc
 from .errors import InterfaceError, OperationalError, ProgrammingError
-from .expr import (ExprParser, build_null_scalar, build_string_scalar,
-                   build_bool_scalar, build_double_scalar, build_int_scalar)
+from .expr import ExprParser, build_expr, build_scalar, build_bool_scalar
 from .helpers import encode_to_bytes, get_item_or_attr
 from .result import ColumnMetaData
 from .protobuf import SERVER_MESSAGES, Message, mysqlxpb_enum
@@ -135,8 +134,7 @@ def get_binding_scalars(self, statement):
                 raise ProgrammingError("Unable to find placeholder for "
                                        "parameter: {0}".format(name))
             pos = statement._binding_map[name]
-            scalars[pos] = self.arg_object_to_scalar(binding["value"],
-                                                     not statement._doc_based)
+            scalars[pos] = build_scalar(binding["value"]).get_message()
         return scalars
 
     def _apply_filter(self, message, statement):
@@ -193,8 +191,7 @@ def send_update(self, stmt):
             operation["operation"] = update_op.update_type
             operation["source"] = update_op.source
             if update_op.value is not None:
-                operation["value"] = self.arg_object_to_expr(
-                    update_op.value, not stmt._doc_based)
+                operation["value"] = build_expr(update_op.value)
             msg["operation"].extend([operation.get_message()])
 
         self._writer.write_message(
@@ -241,13 +238,12 @@ def send_insert(self, stmt):
             row = Message("Mysqlx.Crud.Insert.TypedRow")
             if isinstance(value, list):
                 for val in value:
-                    obj = self.arg_object_to_expr(val, not stmt._doc_based)
-                    row["field"].extend([obj.get_message()])
+                    row["field"].extend([build_expr(val).get_message()])
             else:
-                obj = self.arg_object_to_expr(value, not stmt._doc_based)
-                row["field"].extend([obj.get_message()])
+                row["field"].extend([build_expr(value).get_message()])
             msg["row"].extend([row.get_message()])
 
+        msg["upsert"] = stmt._upsert
         self._writer.write_message(
             mysqlxpb_enum("Mysqlx.ClientMessages.Type.CRUD_INSERT"), msg)
 
@@ -338,39 +334,6 @@ def get_column_metadata(self, rs):
             columns.append(col)
         return columns
 
-    def arg_object_to_expr(self, value, allow_relational):
-        literal = None
-        if value is None:
-            literal = build_null_scalar()
-        if isinstance(value, bool):
-            literal = build_bool_scalar(value)
-        elif isinstance(value, INT_TYPES):
-            literal = build_int_scalar(value)
-        elif isinstance(value, (float)):
-            literal = build_double_scalar(value)
-        elif isinstance(value, STRING_TYPES):
-            try:
-                expression = ExprParser(value, allow_relational).expr()
-                if expression.has_identifier():
-                    msg = Message("Mysqlx.Expr.Expr",
-                                  literal=build_string_scalar(value))
-                    return msg
-                return expression.serialize_to_string()
-            except:
-                literal = build_string_scalar(value)
-        elif isinstance(value, DbDoc):
-            literal = build_string_scalar(str(value))
-        if literal is None:
-            raise InterfaceError("Unsupported type: {0}".format(type(value)))
-
-        msg = Message("Mysqlx.Expr.Expr")
-        msg["type"] = mysqlxpb_enum("Mysqlx.Expr.Expr.Type.LITERAL")
-        msg["literal"] = literal
-        return msg
-
-    def arg_object_to_scalar(self, value, allow_relational):
-        return self.arg_object_to_expr(value, allow_relational).literal
-
     def read_ok(self):
         msg = self._reader.read_message()
         if msg.type == "Mysqlx.Error":

lib/mysqlx/statement.py

@@ -306,15 +306,34 @@ def execute(self):
         return SqlResult(self._connection)
 
 
-class AddStatement(Statement):
+class WriteStatement(Statement):
+    """Provide common write operation attributes
+    """
+    def __init__(self, target, doc_based):
+        super(WriteStatement, self).__init__(target, doc_based)
+        self._values = []
+        self._upsert = False
+
+    def upsert(self, val=True):
+        """Sets the upset flag to the boolean of the value provided.
+        Setting of this flag allows updating of the matched rows/documents
+        with the provided value.
+
+        Args:
+            val (optional[bool]): Set or unset the upsert flag.
+        """
+        self._upsert = val
+        return self
+
+
+class AddStatement(WriteStatement):
     """A statement for document addition on a collection.
 
     Args:
         collection (mysqlx.Collection): The Collection object.
     """
     def __init__(self, collection):
-        super(AddStatement, self).__init__(target=collection)
-        self._values = []
+        super(AddStatement, self).__init__(collection, True)
         self._ids = []
 
     def add(self, *values):
@@ -610,18 +629,16 @@ def get_sql(self):
 
         return stmt
 
-
-class InsertStatement(Statement):
+class InsertStatement(WriteStatement):
     """A statement for insert operations on Table.
 
     Args:
         table (mysqlx.Table): The Table object.
         *fields: The fields to be inserted.
     """
     def __init__(self, table, *fields):
-        super(InsertStatement, self).__init__(target=table, doc_based=False)
+        super(InsertStatement, self).__init__(table, False)
         self._fields = flexible_params(*fields)
-        self._values = []
 
     def values(self, *values):
         """Set the values to be inserted.

src/mysqlxpb/mysqlx/protocol/mysqlx_crud.proto

@@ -150,6 +150,7 @@ message Find {
 // :param projection: name of the columns to insert data into (empty if data_model is DOCUMENT)
 // :param row: set of rows to insert into the collection/table (a single expression with a JSON document literal or an OBJECT expression)
 // :param args: values for parameters used in row expressions
+// :param upsert: true if this should be treated as an Upsert (that is, update on duplicate key)
 // :Returns: :protobuf:msg:`Mysqlx.Resultset::`
 message Insert {
   required Collection collection = 1;
@@ -161,7 +162,8 @@ message Insert {
     repeated Mysqlx.Expr.Expr field = 1;
   };
   repeated TypedRow row = 4;
-  repeated Mysqlx.Datatypes.Scalar args = 5; 
+  repeated Mysqlx.Datatypes.Scalar args = 5;
+  optional bool upsert = 6 [default = false];
 };
 
 // Update documents/rows in a collection/table