M sqlhelp/__init__.py +7 -2
@@ 32,17 32,22 @@ class insert(_sqlbase):
usage: insert('mytable').values(v1=42, v2='Babar')
"""
- __slots__ = ('_table', '_kw')
+ __slots__ = ('_table', '_kw', '_returning')
def __init__(self, table):
self._table = table
self._kw = None
+ self._returning = 'id'
def values(self, **values):
"""Declare values to insert."""
self._kw = values
return self
+ def returning(self, name):
+ self._returning = name
+ return self
+
def _assemble(self):
"""Return the sql query string and values suitable for executor."""
assert self._kw, "Call values() first"
@@ 51,7 56,7 @@ class insert(_sqlbase):
return (
f'insert into {self._table} '
f'({names}) values ({holders}) '
- 'returning id'
+ f'returning {self._returning}'
), self._kw.copy()
M test/test_sqlhelp.py +14 -0
@@ 40,6 40,20 @@ def test_insert():
"{'a': 1}]"
)
+ q = insert(
+ 'table1'
+ ).values(
+ a=1,
+ b=2
+ ).returning('b')
+
+ assert str(q) == (
+ "query::[insert into table1 (a,b) "
+ "values (%(a)s,%(b)s) "
+ "returning b "
+ "{'a': 1, 'b': 2}]"
+ )
+
def test_update():
q = update(