forked from platypush/platypush
Improved support for bulk database statements
- Wrapped insert/update/delete operations in transactions - Proper (and much more efficient) bulk logic - Better upsert logic - Return inserted/updated records if the engine supports it
This commit is contained in:
parent
a90aa2cb2e
commit
0143dac216
1 changed files with 108 additions and 37 deletions
|
@ -7,7 +7,8 @@ from typing import Optional
|
||||||
|
|
||||||
from sqlalchemy import create_engine, Table, MetaData
|
from sqlalchemy import create_engine, Table, MetaData
|
||||||
from sqlalchemy.engine import Engine
|
from sqlalchemy.engine import Engine
|
||||||
from sqlalchemy.sql import text
|
from sqlalchemy.exc import CompileError
|
||||||
|
from sqlalchemy.sql import and_, or_, text
|
||||||
|
|
||||||
from platypush.plugins import Plugin, action
|
from platypush.plugins import Plugin, action
|
||||||
|
|
||||||
|
@ -251,7 +252,9 @@ class DbPlugin(Plugin):
|
||||||
:type key_columns: list
|
:type key_columns: list
|
||||||
:param on_duplicate_update: If set, update the records in case of
|
:param on_duplicate_update: If set, update the records in case of
|
||||||
duplicate rows (default: False). If set, you'll need to specify
|
duplicate rows (default: False). If set, you'll need to specify
|
||||||
``key_columns`` as well.
|
``key_columns`` as well. If ``key_columns`` is set, existing
|
||||||
|
records are found but ``on_duplicate_update`` is false, then
|
||||||
|
existing records will be ignored.
|
||||||
:type on_duplicate_update: bool
|
:type on_duplicate_update: bool
|
||||||
:param args: Extra arguments that will be passed to
|
:param args: Extra arguments that will be passed to
|
||||||
``sqlalchemy.create_engine`` (see
|
``sqlalchemy.create_engine`` (see
|
||||||
|
@ -260,6 +263,9 @@ class DbPlugin(Plugin):
|
||||||
``sqlalchemy.create_engine``
|
``sqlalchemy.create_engine``
|
||||||
(see https:///docs.sqlalchemy.org/en/latest/core/engines.html)
|
(see https:///docs.sqlalchemy.org/en/latest/core/engines.html)
|
||||||
|
|
||||||
|
:return: The inserted records, if the underlying engine supports the
|
||||||
|
``RETURNING`` statement, otherwise nothing.
|
||||||
|
|
||||||
Example:
|
Example:
|
||||||
|
|
||||||
Request::
|
Request::
|
||||||
|
@ -290,26 +296,98 @@ class DbPlugin(Plugin):
|
||||||
key_columns = []
|
key_columns = []
|
||||||
|
|
||||||
engine = self._get_engine(engine, *args, **kwargs)
|
engine = self._get_engine(engine, *args, **kwargs)
|
||||||
|
table, engine = self._get_table(table, engine=engine, *args, **kwargs)
|
||||||
|
insert_records = records
|
||||||
|
update_records = []
|
||||||
|
returned_records = []
|
||||||
|
|
||||||
|
with engine.connect() as connection:
|
||||||
|
# Upsert case
|
||||||
|
if key_columns:
|
||||||
|
insert_records, update_records = self._get_new_and_existing_records(
|
||||||
|
connection, table, records, key_columns
|
||||||
|
)
|
||||||
|
|
||||||
|
with connection.begin():
|
||||||
|
if insert_records:
|
||||||
|
insert = table.insert().values(insert_records)
|
||||||
|
ret = self._execute_try_returning(connection, insert)
|
||||||
|
if ret:
|
||||||
|
returned_records += ret
|
||||||
|
|
||||||
|
if update_records and on_duplicate_update:
|
||||||
|
ret = self._update(connection, table, update_records, key_columns)
|
||||||
|
if ret:
|
||||||
|
returned_records = ret + returned_records
|
||||||
|
|
||||||
|
if returned_records:
|
||||||
|
return returned_records
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _execute_try_returning(connection, stmt):
|
||||||
|
ret = None
|
||||||
|
stmt_with_ret = stmt.returning('*')
|
||||||
|
|
||||||
|
try:
|
||||||
|
ret = connection.execute(stmt_with_ret)
|
||||||
|
except CompileError as e:
|
||||||
|
if str(e).startswith('RETURNING is not supported'):
|
||||||
|
connection.execute(stmt)
|
||||||
|
else:
|
||||||
|
raise e
|
||||||
|
|
||||||
|
if ret:
|
||||||
|
return [
|
||||||
|
{col.name: getattr(row, col.name, None) for col in stmt.table.c}
|
||||||
|
for row in ret
|
||||||
|
]
|
||||||
|
|
||||||
|
def _get_new_and_existing_records(self, connection, table, records, key_columns):
|
||||||
|
records_by_key = {
|
||||||
|
tuple(record.get(k) for k in key_columns): record for record in records
|
||||||
|
}
|
||||||
|
|
||||||
|
query = table.select().where(
|
||||||
|
or_(
|
||||||
|
and_(
|
||||||
|
self._build_condition(table, k, record.get(k)) for k in key_columns
|
||||||
|
)
|
||||||
|
for record in records
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
existing_records = {
|
||||||
|
tuple(getattr(record, k, None) for k in key_columns): record
|
||||||
|
for record in connection.execute(query).all()
|
||||||
|
}
|
||||||
|
|
||||||
|
update_records = [
|
||||||
|
record for k, record in records_by_key.items() if k in existing_records
|
||||||
|
]
|
||||||
|
|
||||||
|
insert_records = [
|
||||||
|
record for k, record in records_by_key.items() if k not in existing_records
|
||||||
|
]
|
||||||
|
|
||||||
|
return insert_records, update_records
|
||||||
|
|
||||||
|
def _update(self, connection, table, records, key_columns):
|
||||||
|
updated_records = []
|
||||||
for record in records:
|
for record in records:
|
||||||
table, engine = self._get_table(table, engine=engine, *args, **kwargs)
|
key = {k: v for (k, v) in record.items() if k in key_columns}
|
||||||
|
values = {k: v for (k, v) in record.items() if k not in key_columns}
|
||||||
|
update = table.update()
|
||||||
|
|
||||||
insert = table.insert().values(**record)
|
for (k, v) in key.items():
|
||||||
|
update = update.where(self._build_condition(table, k, v))
|
||||||
|
|
||||||
try:
|
update = update.values(**values)
|
||||||
engine.execute(insert)
|
ret = self._execute_try_returning(connection, update)
|
||||||
except Exception as e:
|
if ret:
|
||||||
if on_duplicate_update and key_columns:
|
updated_records += ret
|
||||||
self.update(
|
|
||||||
table=table,
|
if updated_records:
|
||||||
records=records,
|
return updated_records
|
||||||
key_columns=key_columns,
|
|
||||||
engine=engine,
|
|
||||||
*args,
|
|
||||||
**kwargs
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
raise e
|
|
||||||
|
|
||||||
@action
|
@action
|
||||||
def update(self, table, records, key_columns, engine=None, *args, **kwargs):
|
def update(self, table, records, key_columns, engine=None, *args, **kwargs):
|
||||||
|
@ -331,6 +409,9 @@ class DbPlugin(Plugin):
|
||||||
``sqlalchemy.create_engine``
|
``sqlalchemy.create_engine``
|
||||||
(see https:///docs.sqlalchemy.org/en/latest/core/engines.html)
|
(see https:///docs.sqlalchemy.org/en/latest/core/engines.html)
|
||||||
|
|
||||||
|
:return: The inserted records, if the underlying engine supports the
|
||||||
|
``RETURNING`` statement, otherwise nothing.
|
||||||
|
|
||||||
Example:
|
Example:
|
||||||
|
|
||||||
Request::
|
Request::
|
||||||
|
@ -357,21 +438,10 @@ class DbPlugin(Plugin):
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
"""
|
"""
|
||||||
|
|
||||||
engine = self._get_engine(engine, *args, **kwargs)
|
engine = self._get_engine(engine, *args, **kwargs)
|
||||||
|
with engine.connect() as connection:
|
||||||
for record in records:
|
|
||||||
table, engine = self._get_table(table, engine=engine, *args, **kwargs)
|
table, engine = self._get_table(table, engine=engine, *args, **kwargs)
|
||||||
key = {k: v for (k, v) in record.items() if k in key_columns}
|
return self._update(connection, table, records, key_columns)
|
||||||
values = {k: v for (k, v) in record.items() if k not in key_columns}
|
|
||||||
|
|
||||||
update = table.update()
|
|
||||||
|
|
||||||
for (k, v) in key.items():
|
|
||||||
update = update.where(self._build_condition(table, k, v))
|
|
||||||
|
|
||||||
update = update.values(**values)
|
|
||||||
engine.execute(update)
|
|
||||||
|
|
||||||
@action
|
@action
|
||||||
def delete(self, table, records, engine=None, *args, **kwargs):
|
def delete(self, table, records, engine=None, *args, **kwargs):
|
||||||
|
@ -412,14 +482,15 @@ class DbPlugin(Plugin):
|
||||||
|
|
||||||
engine = self._get_engine(engine, *args, **kwargs)
|
engine = self._get_engine(engine, *args, **kwargs)
|
||||||
|
|
||||||
for record in records:
|
with engine.connect() as connection, connection.begin():
|
||||||
table, engine = self._get_table(table, engine=engine, *args, **kwargs)
|
for record in records:
|
||||||
delete = table.delete()
|
table, engine = self._get_table(table, engine=engine, *args, **kwargs)
|
||||||
|
delete = table.delete()
|
||||||
|
|
||||||
for (k, v) in record.items():
|
for (k, v) in record.items():
|
||||||
delete = delete.where(self._build_condition(table, k, v))
|
delete = delete.where(self._build_condition(table, k, v))
|
||||||
|
|
||||||
engine.execute(delete)
|
connection.execute(delete)
|
||||||
|
|
||||||
|
|
||||||
# vim:sw=4:ts=4:et:
|
# vim:sw=4:ts=4:et:
|
||||||
|
|
Loading…
Reference in a new issue