Skip to content

Commit

Permalink
Merge pull request pudo#169 from saun4app/feature_insert_ignore
Browse files Browse the repository at this point in the history
Added insert_ignore and changed upsert.
  • Loading branch information
pudo committed May 14, 2016
2 parents fb48c91 + 72b01df commit 5f1a38a
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 27 deletions.
82 changes: 56 additions & 26 deletions dataset/persistence/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,31 @@ def insert(self, row, ensure=None, types={}):
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.
Expand Down Expand Up @@ -156,6 +181,27 @@ def update(self, row, keys, ensure=None, types={}):
except KeyError:
return 0

def _upsert_pre_check(self, row, keys, ensure):
# check whether keys arg is a string and format as a list
try:
if not isinstance(keys, (list, tuple)):
keys = [keys]
self._check_dropped()

ensure = self.database.ensure_schema if ensure is None else ensure
if ensure:
self.create_index(keys)

filters = {}
for key in keys:
filters[key] = row.get(key)

res = self.find_one(**filters)
except:
res = None

return res

def upsert(self, row, keys, ensure=None, types={}):
"""
An UPSERT is a smart combination of insert and update.
Expand All @@ -167,33 +213,17 @@ def upsert(self, row, keys, ensure=None, types={}):
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
if not isinstance(keys, (list, tuple)):
keys = [keys]
self._check_dropped()

ensure = self.database.ensure_schema if ensure is None else ensure
if ensure:
self.create_index(keys)

filters = {}
for key in keys:
filters[key] = row.get(key)

res = self.find_one(**filters)
if res is not None:
row_count = self.update(row, keys, ensure=ensure, types=types)
if row_count == 0:
return False
elif row_count == 1:
try:
return res['id']
except KeyError:
return True
else:
return True
else:
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):
"""
Expand Down
2 changes: 1 addition & 1 deletion docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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__


Expand Down
25 changes: 25 additions & 0 deletions test/test_persistence.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,31 @@ def test_insert(self):
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),
Expand Down

0 comments on commit 5f1a38a

Please sign in to comment.