From 19d6984e727d79af79c466b0c3842de7a4a46207 Mon Sep 17 00:00:00 2001
From: Blake Eggleston <bdeggleston@gmail.com>
Date: Wed, 24 Jul 2013 09:49:15 -0700
Subject: [PATCH] updating the model and query classes to handle models with
 counter columns correctly

---
 cqlengine/columns.py | 26 +++++++++++++++-----------
 cqlengine/models.py  | 10 ++++++++--
 cqlengine/query.py   |  9 +++++----
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/cqlengine/columns.py b/cqlengine/columns.py
index 0f505ccc..56221c2d 100644
--- a/cqlengine/columns.py
+++ b/cqlengine/columns.py
@@ -249,22 +249,26 @@ class Integer(Column):
 class Counter(Integer):
     db_type = 'counter'
 
+    def __init__(self,
+                 index=False,
+                 db_field=None,
+                 required=False):
+        super(Counter, self).__init__(
+            primary_key=False,
+            partition_key=False,
+            index=index,
+            db_field=db_field,
+            default=0,
+            required=required,
+        )
+
     def get_update_statement(self, val, prev, ctx):
         val = self.to_database(val)
-        prev = self.to_database(prev)
+        prev = self.to_database(prev or 0)
         field_id = uuid4().hex
 
-        # use a set statement if there is no
-        # previous value to compute a delta from
-        if prev is None:
-            ctx[field_id] = val
-            return ['"{}" = :{}'.format(self.db_field_name, field_id)]
-
         delta = val - prev
-        if delta == 0:
-            return []
-
-        sign = '+' if delta > 0 else '-'
+        sign = '-' if delta < 0 else '+'
         delta = abs(delta)
         ctx[field_id] = delta
         return ['"{0}" = "{0}" {1} {2}'.format(self.db_field_name, sign, delta)]
diff --git a/cqlengine/models.py b/cqlengine/models.py
index a2d9b6ab..640e023b 100644
--- a/cqlengine/models.py
+++ b/cqlengine/models.py
@@ -70,6 +70,7 @@ class ColumnQueryEvaluator(AbstractQueryableColumn):
     def _get_column(self):
         return self.column
 
+
 class ColumnDescriptor(object):
     """
     Handles the reading and writing of column values to and from
@@ -127,6 +128,7 @@ class BaseModel(object):
     """
 
     class DoesNotExist(_DoesNotExist): pass
+
     class MultipleObjectsReturned(_MultipleObjectsReturned): pass
 
     objects = QuerySetDescriptor()
@@ -316,14 +318,17 @@ class ModelMetaClass(type):
 
         column_definitions = inherited_columns.items() + column_definitions
 
-        #columns defined on model, excludes automatically
-        #defined columns
         defined_columns = OrderedDict(column_definitions)
 
         #prepend primary key if one hasn't been defined
         if not is_abstract and not any([v.primary_key for k,v in column_definitions]):
             raise ModelDefinitionException("At least 1 primary key is required.")
 
+        counter_columns = [c for c in defined_columns.values() if isinstance(c, columns.Counter)]
+        data_columns = [c for c in defined_columns.values() if not c.primary_key and not isinstance(c, columns.Counter)]
+        if counter_columns and data_columns:
+            raise ModelDefinitionException('counter models may not have data columns')
+
         has_partition_keys = any(v.partition_key for (k, v) in column_definitions)
 
         #TODO: check that the defined columns don't conflict with any of the Model API's existing attributes/methods
@@ -382,6 +387,7 @@ class ModelMetaClass(type):
 
         attrs['_partition_keys'] = partition_keys
         attrs['_clustering_keys'] = clustering_keys
+        attrs['_has_counter'] = len(counter_columns) > 0
 
         #setup class exceptions
         DoesNotExistBase = None
diff --git a/cqlengine/query.py b/cqlengine/query.py
index 4327ff11..194b4c93 100644
--- a/cqlengine/query.py
+++ b/cqlengine/query.py
@@ -801,7 +801,7 @@ class DMLQuery(object):
         query_values = {field_ids[n]:field_values[n] for n in field_names}
 
         qs = []
-        if self.instance._can_update():
+        if self.instance._has_counter or self.instance._can_update():
             qs += ["UPDATE {}".format(self.column_family_name)]
             qs += ["SET"]
 
@@ -818,7 +818,7 @@ class DMLQuery(object):
 
                         val_mgr = self.instance._values[name]
                         set_statements += col.get_update_statement(val, val_mgr.previous_value, query_values)
-                        pass
+
                     else:
                         set_statements += ['"{}" = :{}'.format(col.db_field_name, field_ids[col.db_field_name])]
             qs += [', '.join(set_statements)]
@@ -831,8 +831,9 @@ class DMLQuery(object):
 
             qs += [' AND '.join(where_statements)]
 
-            # clear the qs if there are not set statements
-            if not set_statements: qs = []
+            # clear the qs if there are no set statements and this is not a counter model
+            if not set_statements and not self.instance._has_counter:
+                qs = []
 
         else:
             qs += ["INSERT INTO {}".format(self.column_family_name)]