Add map item deletion in cqlengine

aboudreault · aboudreault · commit ae4fbd34a921 · 2017-04-03T15:07:31.000-04:00
diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py
@@ -1252,6 +1252,9 @@ class Row(Model):
 
             # add items to a map
             Row.objects(row_id=5).update(map_column__update={1: 2, 3: 4})
+
+            # remove items from a map
+            Row.objects(row_id=5).update(map_column__remove={1, 2})
         """
         if not values:
             return
@@ -1270,8 +1273,14 @@ class Row(Model):
             if col.is_primary_key:
                 raise ValidationError("Cannot apply update to primary key '{0}' for {1}.{2}".format(col_name, self.__module__, self.model.__name__))
 
-            # we should not provide default values in this use case.
-            val = col.validate(val)
+            if col_op == 'remove' and isinstance(col, columns.Map):
+                if not isinstance(val, set):
+                    raise ValidationError(
+                        "Cannot apply update operation '{0}' on column '{1}' with value '{2}'. A set is required.".format(col_op, col_name, val))
+                val = {v: None for v in val}
+            else:
+                # we should not provide default values in this use case.
+                val = col.validate(val)
 
             if val is None:
                 nulled_columns.add(col_name)
diff --git a/cassandra/cqlengine/statements.py b/cassandra/cqlengine/statements.py
@@ -359,10 +359,13 @@ class MapUpdateClause(ContainerUpdateClause):
     col_type = columns.Map
 
     _updates = None
+    _removals = None
 
     def _analyze(self):
         if self._operation == "update":
             self._updates = self.value.keys()
+        elif self._operation == "remove":
+            self._removals = {v for v in self.value.keys()}
         else:
             if self.previous is None:
                 self._updates = sorted([k for k, v in self.value.items()])
@@ -373,12 +376,14 @@ def _analyze(self):
     def get_context_size(self):
         if self.is_assignment:
             return 1
-        return len(self._updates or []) * 2
+        return int((len(self._updates or []) * 2) + int(bool(self._removals)))
 
     def update_context(self, ctx):
         ctx_id = self.context_id
         if self.is_assignment:
             ctx[str(ctx_id)] = {}
+        elif self._removals is not None:
+            ctx[str(ctx_id)] = self._removals
         else:
             for key in self._updates or []:
                 val = self.value.get(key)
@@ -390,14 +395,17 @@ def update_context(self, ctx):
     def is_assignment(self):
         if not self._analyzed:
             self._analyze()
-        return self.previous is None and not self._updates
+        return self.previous is None and not self._updates and not self._removals
 
     def __unicode__(self):
         qs = []
 
         ctx_id = self.context_id
         if self.is_assignment:
             qs += ['"{0}" = %({1})s'.format(self.field, ctx_id)]
+        elif self._removals is not None:
+            qs += ['"{0}" = "{0}" - %({1})s'.format(self.field, ctx_id)]
+            ctx_id += 1
         else:
             for _ in self._updates or []:
                 qs += ['"{0}"[%({1})s] = %({2})s'.format(self.field, ctx_id, ctx_id + 1)]
