Merge pull request #169 from saun4app/feature_insert_ignore

Added insert_ignore and changed upsert.

dataset/persistence/table.py

@@ -79,6 +79,31 @@ class Table(object):
         if len(res.inserted_primary_key) > 0:
             return res.inserted_primary_key[0]
 
+    def insert_ignore(self, row, keys, ensure=None, types={}):
+        """
+        Add a row (type: dict) into the table if the row does not exist.
+
+        If rows with matching ``keys`` exist they will be added to the table.
+
+        Setting ``ensure`` results in automatically creating missing columns,
+        i.e., keys of the row are not table columns.
+
+        During column creation, ``types`` will be checked for a key
+        matching the name of a column to be created, and the given
+        SQLAlchemy column type will be used. Otherwise, the type is
+        guessed from the row value, defaulting to a simple unicode
+        field.
+        ::
+
+            data = dict(id=10, title='I am a banana!')
+            table.insert_ignore(data, ['id'])
+        """
+        res = self._upsert_pre_check(row, keys, ensure)
+        if res is None:
+            return self.insert(row, ensure=ensure, types=types)
+        else:
+            return False
+
     def insert_many(self, rows, chunk_size=1000, ensure=None, types={}):
         """
         Add many rows at a time.
@@ -156,18 +181,9 @@ class Table(object):
         except KeyError:
             return 0
 
-    def upsert(self, row, keys, ensure=None, types={}):
+    def _upsert_pre_check(self, row, keys, ensure):
-        """
-        An UPSERT is a smart combination of insert and update.
-
-        If rows with matching ``keys`` exist they will be updated, otherwise a
-        new row is inserted in the table.
-        ::
-
-            data = dict(id=10, title='I am a banana!')
-            table.upsert(data, ['id'])
-        """
         # check whether keys arg is a string and format as a list
+        try:
             if not isinstance(keys, (list, tuple)):
                 keys = [keys]
             self._check_dropped()
@@ -181,19 +197,33 @@ class Table(object):
                 filters[key] = row.get(key)
 
             res = self.find_one(**filters)
-        if res is not None:
+        except:
-            row_count = self.update(row, keys, ensure=ensure, types=types)
+            res = None
-            if row_count == 0:
+
-                return False
+        return res
-            elif row_count == 1:
+
-                try:
+    def upsert(self, row, keys, ensure=None, types={}):
-                    return res['id']
+        """
-                except KeyError:
+        An UPSERT is a smart combination of insert and update.
-                    return True
+
-            else:
+        If rows with matching ``keys`` exist they will be updated, otherwise a
-                return True
+        new row is inserted in the table.
-        else:
+        ::
+
+            data = dict(id=10, title='I am a banana!')
+            table.upsert(data, ['id'])
+        """
+        res = self._upsert_pre_check(row, keys, ensure)
+        if res is None:
             return self.insert(row, ensure=ensure, types=types)
+        else:
+            row_count = self.update(row, keys, ensure=ensure, types=types)
+            try:
+                result = (row_count > 0, res['id'])[row_count == 1]
+            except KeyError:
+                result = row_count > 0
+
+            return result
 
     def delete(self, *_clauses, **_filter):
         """

docs/api.rst

@@ -16,7 +16,7 @@ Table
 -----
 
 .. autoclass:: dataset.Table
-   :members: columns, find, find_one, all, count, distinct, insert, insert_many, update, upsert, delete, create_column, drop_column, create_index, drop
+   :members: columns, find, find_one, all, count, distinct, insert, insert_ignore, insert_many, update, upsert, delete, create_column, drop_column, create_index, drop
    :special-members: __len__, __iter__
 
 

test/test_persistence.py

@@ -170,6 +170,31 @@ class TableTestCase(unittest.TestCase):
         assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
         assert self.tbl.find_one(id=last_id)['place'] == 'Berlin'
 
+    def test_insert_ignore(self):
+        self.tbl.insert_ignore({
+            'date': datetime(2011, 1, 2),
+            'temperature': -10,
+            'place': 'Berlin'},
+            ['place']
+        )
+        assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
+        self.tbl.insert_ignore({
+            'date': datetime(2011, 1, 2),
+            'temperature': -10,
+            'place': 'Berlin'},
+            ['place']
+        )
+        assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
+
+    def test_insert_ignore_all_key(self):
+        for i in range(0, 2):
+            self.tbl.insert_ignore({
+                'date': datetime(2011, 1, 2),
+                'temperature': -10,
+                'place': 'Berlin'},
+                ['date', 'temperature', 'place']
+            )
+
     def test_upsert(self):
         self.tbl.upsert({
             'date': datetime(2011, 1, 2),