Optimise _update_client_ips_batch_txn to batch together database operations. (#12252)

Co-authored-by: Patrick Cloke <clokep@users.noreply.github.com>
This commit is contained in:
reivilibre 2022-04-08 15:29:13 +01:00 committed by GitHub
parent 0cd182f296
commit e630722f11
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 191 additions and 52 deletions

View file

@ -1268,6 +1268,7 @@ class DatabasePool:
value_names: Collection[str],
value_values: Collection[Collection[Any]],
desc: str,
lock: bool = True,
) -> None:
"""
Upsert, many times.
@ -1279,6 +1280,8 @@ class DatabasePool:
value_names: The value column names
value_values: A list of each row's value column values.
Ignored if value_names is empty.
lock: True to lock the table when doing the upsert. Unused if the database engine
supports native upserts.
"""
# We can autocommit if we are going to use native upserts
@ -1286,7 +1289,7 @@ class DatabasePool:
self.engine.can_native_upsert and table not in self._unsafe_to_upsert_tables
)
return await self.runInteraction(
await self.runInteraction(
desc,
self.simple_upsert_many_txn,
table,
@ -1294,6 +1297,7 @@ class DatabasePool:
key_values,
value_names,
value_values,
lock=lock,
db_autocommit=autocommit,
)
@ -1305,6 +1309,7 @@ class DatabasePool:
key_values: Collection[Iterable[Any]],
value_names: Collection[str],
value_values: Iterable[Iterable[Any]],
lock: bool = True,
) -> None:
"""
Upsert, many times.
@ -1316,6 +1321,8 @@ class DatabasePool:
value_names: The value column names
value_values: A list of each row's value column values.
Ignored if value_names is empty.
lock: True to lock the table when doing the upsert. Unused if the database engine
supports native upserts.
"""
if self.engine.can_native_upsert and table not in self._unsafe_to_upsert_tables:
return self.simple_upsert_many_txn_native_upsert(
@ -1323,7 +1330,7 @@ class DatabasePool:
)
else:
return self.simple_upsert_many_txn_emulated(
txn, table, key_names, key_values, value_names, value_values
txn, table, key_names, key_values, value_names, value_values, lock=lock
)
def simple_upsert_many_txn_emulated(
@ -1334,6 +1341,7 @@ class DatabasePool:
key_values: Collection[Iterable[Any]],
value_names: Collection[str],
value_values: Iterable[Iterable[Any]],
lock: bool = True,
) -> None:
"""
Upsert, many times, but without native UPSERT support or batching.
@ -1345,17 +1353,24 @@ class DatabasePool:
value_names: The value column names
value_values: A list of each row's value column values.
Ignored if value_names is empty.
lock: True to lock the table when doing the upsert.
"""
# No value columns, therefore make a blank list so that the following
# zip() works correctly.
if not value_names:
value_values = [() for x in range(len(key_values))]
if lock:
# Lock the table just once, to prevent it being done once per row.
# Note that, according to Postgres' documentation, once obtained,
# the lock is held for the remainder of the current transaction.
self.engine.lock_table(txn, "user_ips")
for keyv, valv in zip(key_values, value_values):
_keys = {x: y for x, y in zip(key_names, keyv)}
_vals = {x: y for x, y in zip(value_names, valv)}
self.simple_upsert_txn_emulated(txn, table, _keys, _vals)
self.simple_upsert_txn_emulated(txn, table, _keys, _vals, lock=False)
def simple_upsert_many_txn_native_upsert(
self,
@ -1792,6 +1807,86 @@ class DatabasePool:
return txn.rowcount
async def simple_update_many(
self,
table: str,
key_names: Collection[str],
key_values: Collection[Iterable[Any]],
value_names: Collection[str],
value_values: Iterable[Iterable[Any]],
desc: str,
) -> None:
"""
Update, many times, using batching where possible.
If the keys don't match anything, nothing will be updated.
Args:
table: The table to update
key_names: The key column names.
key_values: A list of each row's key column values.
value_names: The names of value columns to update.
value_values: A list of each row's value column values.
"""
await self.runInteraction(
desc,
self.simple_update_many_txn,
table,
key_names,
key_values,
value_names,
value_values,
)
@staticmethod
def simple_update_many_txn(
txn: LoggingTransaction,
table: str,
key_names: Collection[str],
key_values: Collection[Iterable[Any]],
value_names: Collection[str],
value_values: Collection[Iterable[Any]],
) -> None:
"""
Update, many times, using batching where possible.
If the keys don't match anything, nothing will be updated.
Args:
table: The table to update
key_names: The key column names.
key_values: A list of each row's key column values.
value_names: The names of value columns to update.
value_values: A list of each row's value column values.
"""
if len(value_values) != len(key_values):
raise ValueError(
f"{len(key_values)} key rows and {len(value_values)} value rows: should be the same number."
)
# List of tuples of (value values, then key values)
# (This matches the order needed for the query)
args = [tuple(x) + tuple(y) for x, y in zip(value_values, key_values)]
for ks, vs in zip(key_values, value_values):
args.append(tuple(vs) + tuple(ks))
# 'col1 = ?, col2 = ?, ...'
set_clause = ", ".join(f"{n} = ?" for n in value_names)
if key_names:
# 'WHERE col3 = ? AND col4 = ? AND col5 = ?'
where_clause = "WHERE " + (" AND ".join(f"{n} = ?" for n in key_names))
else:
where_clause = ""
# UPDATE mytable SET col1 = ?, col2 = ? WHERE col3 = ? AND col4 = ?
sql = f"""
UPDATE {table} SET {set_clause} {where_clause}
"""
txn.execute_batch(sql, args)
async def simple_update_one(
self,
table: str,