2019-12-04 08:52:46 -05:00
|
|
|
#
|
2023-11-21 15:29:58 -05:00
|
|
|
# This file is licensed under the Affero General Public License (AGPL) version 3.
|
|
|
|
#
|
2024-01-23 06:26:48 -05:00
|
|
|
# Copyright 2019 The Matrix.org Foundation C.I.C.
|
|
|
|
# Copyright 2014-2016 OpenMarket Ltd
|
2023-11-21 15:29:58 -05:00
|
|
|
# Copyright (C) 2023 New Vector, Ltd
|
|
|
|
#
|
|
|
|
# This program is free software: you can redistribute it and/or modify
|
|
|
|
# it under the terms of the GNU Affero General Public License as
|
|
|
|
# published by the Free Software Foundation, either version 3 of the
|
|
|
|
# License, or (at your option) any later version.
|
|
|
|
#
|
|
|
|
# See the GNU Affero General Public License for more details:
|
|
|
|
# <https://www.gnu.org/licenses/agpl-3.0.html>.
|
|
|
|
#
|
|
|
|
# Originally licensed under the Apache License, Version 2.0:
|
|
|
|
# <http://www.apache.org/licenses/LICENSE-2.0>.
|
|
|
|
#
|
|
|
|
# [This file includes modifications made by New Vector Limited]
|
2019-12-04 08:52:46 -05:00
|
|
|
#
|
|
|
|
#
|
2021-12-13 14:01:27 -05:00
|
|
|
import inspect
|
2019-12-04 08:52:46 -05:00
|
|
|
import logging
|
|
|
|
import time
|
2021-12-13 14:01:27 -05:00
|
|
|
import types
|
2021-08-02 09:24:43 -04:00
|
|
|
from collections import defaultdict
|
2020-02-27 06:53:40 -05:00
|
|
|
from time import monotonic as monotonic_time
|
2020-04-07 18:06:39 -04:00
|
|
|
from typing import (
|
2021-10-22 13:15:41 -04:00
|
|
|
TYPE_CHECKING,
|
2020-04-07 18:06:39 -04:00
|
|
|
Any,
|
2022-07-19 07:25:29 -04:00
|
|
|
Awaitable,
|
2020-04-07 18:06:39 -04:00
|
|
|
Callable,
|
2021-04-22 11:43:50 -04:00
|
|
|
Collection,
|
2020-04-07 18:06:39 -04:00
|
|
|
Dict,
|
|
|
|
Iterable,
|
|
|
|
Iterator,
|
|
|
|
List,
|
|
|
|
Optional,
|
2023-08-29 10:38:56 -04:00
|
|
|
Sequence,
|
2020-04-07 18:06:39 -04:00
|
|
|
Tuple,
|
2022-05-12 10:33:50 -04:00
|
|
|
Type,
|
2020-04-07 18:06:39 -04:00
|
|
|
TypeVar,
|
2020-09-02 13:11:02 -04:00
|
|
|
cast,
|
2020-08-26 07:19:32 -04:00
|
|
|
overload,
|
2020-04-07 18:06:39 -04:00
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-10-02 10:20:45 -04:00
|
|
|
import attr
|
2022-07-21 08:23:05 -04:00
|
|
|
from prometheus_client import Counter, Histogram
|
2022-05-09 05:28:38 -04:00
|
|
|
from typing_extensions import Concatenate, Literal, ParamSpec
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2019-12-18 05:45:12 -05:00
|
|
|
from twisted.enterprise import adbapi
|
2022-05-12 10:33:50 -04:00
|
|
|
from twisted.internet.interfaces import IReactorCore
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
from synapse.api.errors import StoreError
|
2019-12-18 05:45:12 -05:00
|
|
|
from synapse.config.database import DatabaseConnectionConfig
|
2021-06-03 11:31:56 -04:00
|
|
|
from synapse.logging import opentracing
|
2020-03-10 10:40:28 -04:00
|
|
|
from synapse.logging.context import (
|
|
|
|
LoggingContext,
|
2020-03-24 10:45:33 -04:00
|
|
|
current_context,
|
2020-03-10 10:40:28 -04:00
|
|
|
make_deferred_yieldable,
|
|
|
|
)
|
2023-06-07 13:12:23 -04:00
|
|
|
from synapse.metrics import LaterGauge, register_threadpool
|
2019-12-04 08:52:46 -05:00
|
|
|
from synapse.metrics.background_process_metrics import run_as_background_process
|
2019-12-04 10:09:36 -05:00
|
|
|
from synapse.storage.background_updates import BackgroundUpdater
|
2020-02-27 06:53:40 -05:00
|
|
|
from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, Sqlite3Engine
|
2023-04-14 14:04:49 -04:00
|
|
|
from synapse.storage.types import Connection, Cursor, SQLQueryParameters
|
2022-07-19 07:25:29 -04:00
|
|
|
from synapse.util.async_helpers import delay_cancellation
|
2021-12-15 12:00:50 -05:00
|
|
|
from synapse.util.iterutils import batch_iter
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2021-10-22 13:15:41 -04:00
|
|
|
if TYPE_CHECKING:
|
|
|
|
from synapse.server import HomeServer
|
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
# python 3 does not have a maximum int value
|
2022-03-29 06:41:19 -04:00
|
|
|
MAX_TXN_ID = 2**63 - 1
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-07-28 13:52:13 -04:00
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
sql_logger = logging.getLogger("synapse.storage.SQL")
|
|
|
|
transaction_logger = logging.getLogger("synapse.storage.txn")
|
|
|
|
perf_logger = logging.getLogger("synapse.storage.TIME")
|
|
|
|
|
|
|
|
sql_scheduling_timer = Histogram("synapse_storage_schedule_time", "sec")
|
|
|
|
|
|
|
|
sql_query_timer = Histogram("synapse_storage_query_time", "sec", ["verb"])
|
2022-07-21 08:23:05 -04:00
|
|
|
sql_txn_count = Counter("synapse_storage_transaction_time_count", "sec", ["desc"])
|
|
|
|
sql_txn_duration = Counter("synapse_storage_transaction_time_sum", "sec", ["desc"])
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
|
|
|
|
# Unique indexes which have been added in background updates. Maps from table name
|
|
|
|
# to the name of the background update which added the unique index to that table.
|
|
|
|
#
|
|
|
|
# This is used by the upsert logic to figure out which tables are safe to do a proper
|
|
|
|
# UPSERT on: until the relevant background update has completed, we
|
|
|
|
# have to emulate an upsert by locking the table.
|
|
|
|
#
|
|
|
|
UNIQUE_INDEX_BACKGROUND_UPDATES = {
|
|
|
|
"user_ips": "user_ips_device_unique_index",
|
|
|
|
"device_lists_remote_extremeties": "device_lists_remote_extremeties_unique_idx",
|
|
|
|
"device_lists_remote_cache": "device_lists_remote_cache_unique_idx",
|
|
|
|
"event_search": "event_search_event_id_idx",
|
2022-05-23 05:28:56 -04:00
|
|
|
"local_media_repository_thumbnails": "local_media_repository_thumbnails_method_idx",
|
|
|
|
"remote_media_cache_thumbnails": "remote_media_repository_thumbnails_method_idx",
|
2022-10-04 09:47:04 -04:00
|
|
|
"event_push_summary": "event_push_summary_unique_index2",
|
2022-09-23 10:33:28 -04:00
|
|
|
"receipts_linearized": "receipts_linearized_unique_index",
|
|
|
|
"receipts_graph": "receipts_graph_unique_index",
|
2019-12-04 08:52:46 -05:00
|
|
|
}
|
|
|
|
|
|
|
|
|
2023-03-09 07:10:09 -05:00
|
|
|
class _PoolConnection(Connection):
|
|
|
|
"""
|
|
|
|
A Connection from twisted.enterprise.adbapi.Connection.
|
|
|
|
"""
|
|
|
|
|
2024-03-13 12:46:44 -04:00
|
|
|
def reconnect(self) -> None: ...
|
2023-03-09 07:10:09 -05:00
|
|
|
|
|
|
|
|
2019-12-18 05:45:12 -05:00
|
|
|
def make_pool(
|
2022-05-12 10:33:50 -04:00
|
|
|
reactor: IReactorCore,
|
|
|
|
db_config: DatabaseConnectionConfig,
|
|
|
|
engine: BaseDatabaseEngine,
|
2019-12-18 05:45:12 -05:00
|
|
|
) -> adbapi.ConnectionPool:
|
2021-02-16 17:32:34 -05:00
|
|
|
"""Get the connection pool for the database."""
|
2019-12-18 05:45:12 -05:00
|
|
|
|
2020-11-12 09:26:24 -05:00
|
|
|
# By default enable `cp_reconnect`. We need to fiddle with db_args in case
|
|
|
|
# someone has explicitly set `cp_reconnect`.
|
|
|
|
db_args = dict(db_config.config.get("args", {}))
|
|
|
|
db_args.setdefault("cp_reconnect", True)
|
|
|
|
|
2022-05-12 10:33:50 -04:00
|
|
|
def _on_new_connection(conn: Connection) -> None:
|
2021-06-08 08:49:29 -04:00
|
|
|
# Ensure we have a logging context so we can correctly track queries,
|
|
|
|
# etc.
|
|
|
|
with LoggingContext("db.on_new_connection"):
|
|
|
|
engine.on_new_connection(
|
|
|
|
LoggingDatabaseConnection(conn, engine, "on_new_connection")
|
|
|
|
)
|
|
|
|
|
2021-11-01 07:21:36 -04:00
|
|
|
connection_pool = adbapi.ConnectionPool(
|
2019-12-18 05:45:12 -05:00
|
|
|
db_config.config["name"],
|
|
|
|
cp_reactor=reactor,
|
2021-06-08 08:49:29 -04:00
|
|
|
cp_openfun=_on_new_connection,
|
2020-11-12 09:26:24 -05:00
|
|
|
**db_args,
|
2019-12-18 05:45:12 -05:00
|
|
|
)
|
|
|
|
|
2021-11-01 07:21:36 -04:00
|
|
|
register_threadpool(f"database-{db_config.name}", connection_pool.threadpool)
|
|
|
|
|
|
|
|
return connection_pool
|
|
|
|
|
2019-12-18 05:45:12 -05:00
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
def make_conn(
|
2020-10-02 10:20:45 -04:00
|
|
|
db_config: DatabaseConnectionConfig,
|
|
|
|
engine: BaseDatabaseEngine,
|
|
|
|
default_txn_name: str,
|
2021-07-06 08:03:16 -04:00
|
|
|
) -> "LoggingDatabaseConnection":
|
2019-12-18 05:45:12 -05:00
|
|
|
"""Make a new connection to the database and return it.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Connection
|
|
|
|
"""
|
|
|
|
|
|
|
|
db_params = {
|
|
|
|
k: v
|
|
|
|
for k, v in db_config.config.get("args", {}).items()
|
|
|
|
if not k.startswith("cp_")
|
|
|
|
}
|
2020-10-02 10:20:45 -04:00
|
|
|
native_db_conn = engine.module.connect(**db_params)
|
|
|
|
db_conn = LoggingDatabaseConnection(native_db_conn, engine, default_txn_name)
|
|
|
|
|
2019-12-18 05:45:12 -05:00
|
|
|
engine.on_new_connection(db_conn)
|
|
|
|
return db_conn
|
|
|
|
|
|
|
|
|
2022-01-13 08:49:28 -05:00
|
|
|
@attr.s(slots=True, auto_attribs=True)
|
2020-10-02 10:20:45 -04:00
|
|
|
class LoggingDatabaseConnection:
|
|
|
|
"""A wrapper around a database connection that returns `LoggingTransaction`
|
|
|
|
as its cursor class.
|
|
|
|
|
|
|
|
This is mainly used on startup to ensure that queries get logged correctly
|
|
|
|
"""
|
|
|
|
|
2022-01-13 08:49:28 -05:00
|
|
|
conn: Connection
|
|
|
|
engine: BaseDatabaseEngine
|
|
|
|
default_txn_name: str
|
2020-10-02 10:20:45 -04:00
|
|
|
|
|
|
|
def cursor(
|
2022-05-12 10:33:50 -04:00
|
|
|
self,
|
|
|
|
*,
|
|
|
|
txn_name: Optional[str] = None,
|
|
|
|
after_callbacks: Optional[List["_CallbackListEntry"]] = None,
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks: Optional[List["_AsyncCallbackListEntry"]] = None,
|
2022-05-12 10:33:50 -04:00
|
|
|
exception_callbacks: Optional[List["_CallbackListEntry"]] = None,
|
2020-10-02 10:20:45 -04:00
|
|
|
) -> "LoggingTransaction":
|
|
|
|
if not txn_name:
|
|
|
|
txn_name = self.default_txn_name
|
|
|
|
|
|
|
|
return LoggingTransaction(
|
|
|
|
self.conn.cursor(),
|
|
|
|
name=txn_name,
|
|
|
|
database_engine=self.engine,
|
|
|
|
after_callbacks=after_callbacks,
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks=async_after_callbacks,
|
2020-10-02 10:20:45 -04:00
|
|
|
exception_callbacks=exception_callbacks,
|
|
|
|
)
|
|
|
|
|
|
|
|
def close(self) -> None:
|
|
|
|
self.conn.close()
|
|
|
|
|
|
|
|
def commit(self) -> None:
|
|
|
|
self.conn.commit()
|
|
|
|
|
2021-02-05 15:39:19 -05:00
|
|
|
def rollback(self) -> None:
|
|
|
|
self.conn.rollback()
|
2020-10-02 10:20:45 -04:00
|
|
|
|
2021-12-13 12:05:00 -05:00
|
|
|
def __enter__(self) -> "LoggingDatabaseConnection":
|
2020-10-02 10:20:45 -04:00
|
|
|
self.conn.__enter__()
|
|
|
|
return self
|
|
|
|
|
2022-05-12 10:33:50 -04:00
|
|
|
def __exit__(
|
|
|
|
self,
|
|
|
|
exc_type: Optional[Type[BaseException]],
|
|
|
|
exc_value: Optional[BaseException],
|
|
|
|
traceback: Optional[types.TracebackType],
|
|
|
|
) -> Optional[bool]:
|
2020-10-02 10:20:45 -04:00
|
|
|
return self.conn.__exit__(exc_type, exc_value, traceback)
|
|
|
|
|
|
|
|
# Proxy through any unknown lookups to the DB conn class.
|
2022-05-12 10:33:50 -04:00
|
|
|
def __getattr__(self, name: str) -> Any:
|
2020-10-02 10:20:45 -04:00
|
|
|
return getattr(self.conn, name)
|
|
|
|
|
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
# The type of entry which goes on our after_callbacks and exception_callbacks lists.
|
2022-05-09 06:27:39 -04:00
|
|
|
_CallbackListEntry = Tuple[Callable[..., object], Tuple[object, ...], Dict[str, object]]
|
2022-07-19 07:25:29 -04:00
|
|
|
_AsyncCallbackListEntry = Tuple[
|
|
|
|
Callable[..., Awaitable], Tuple[object, ...], Dict[str, object]
|
|
|
|
]
|
2020-02-27 06:53:40 -05:00
|
|
|
|
2022-05-09 05:28:38 -04:00
|
|
|
P = ParamSpec("P")
|
2021-01-11 11:09:22 -05:00
|
|
|
R = TypeVar("R")
|
|
|
|
|
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
class LoggingTransaction:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""An object that almost-transparently proxies for the 'txn' object
|
|
|
|
passed to the constructor. Adds logging and metrics to the .execute()
|
|
|
|
method.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
txn: The database transaction object to wrap.
|
2020-02-27 06:53:40 -05:00
|
|
|
name: The name of this transactions for logging.
|
|
|
|
database_engine
|
|
|
|
after_callbacks: A list that callbacks will be appended to
|
2019-12-04 08:52:46 -05:00
|
|
|
that have been added by `call_after` which should be run on
|
|
|
|
successful completion of the transaction. None indicates that no
|
|
|
|
callbacks should be allowed to be scheduled to run.
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks: A list that asynchronous callbacks will be appended
|
|
|
|
to by `async_call_after` which should run, before after_callbacks, on
|
|
|
|
successful completion of the transaction. None indicates that no
|
|
|
|
callbacks should be allowed to be scheduled to run.
|
2020-02-27 06:53:40 -05:00
|
|
|
exception_callbacks: A list that callbacks will be appended
|
2019-12-04 08:52:46 -05:00
|
|
|
to that have been added by `call_on_exception` which should be run
|
|
|
|
if transaction ends with an error. None indicates that no callbacks
|
|
|
|
should be allowed to be scheduled to run.
|
|
|
|
"""
|
|
|
|
|
|
|
|
__slots__ = [
|
|
|
|
"txn",
|
|
|
|
"name",
|
|
|
|
"database_engine",
|
|
|
|
"after_callbacks",
|
2022-07-19 07:25:29 -04:00
|
|
|
"async_after_callbacks",
|
2019-12-04 08:52:46 -05:00
|
|
|
"exception_callbacks",
|
|
|
|
]
|
|
|
|
|
|
|
|
def __init__(
|
2020-02-27 06:53:40 -05:00
|
|
|
self,
|
|
|
|
txn: Cursor,
|
|
|
|
name: str,
|
|
|
|
database_engine: BaseDatabaseEngine,
|
|
|
|
after_callbacks: Optional[List[_CallbackListEntry]] = None,
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks: Optional[List[_AsyncCallbackListEntry]] = None,
|
2020-02-27 06:53:40 -05:00
|
|
|
exception_callbacks: Optional[List[_CallbackListEntry]] = None,
|
2019-12-04 08:52:46 -05:00
|
|
|
):
|
2020-02-27 06:53:40 -05:00
|
|
|
self.txn = txn
|
|
|
|
self.name = name
|
|
|
|
self.database_engine = database_engine
|
|
|
|
self.after_callbacks = after_callbacks
|
2022-07-19 07:25:29 -04:00
|
|
|
self.async_after_callbacks = async_after_callbacks
|
2020-02-27 06:53:40 -05:00
|
|
|
self.exception_callbacks = exception_callbacks
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-05-09 06:27:39 -04:00
|
|
|
def call_after(
|
|
|
|
self, callback: Callable[P, object], *args: P.args, **kwargs: P.kwargs
|
|
|
|
) -> None:
|
2022-03-29 07:31:05 -04:00
|
|
|
"""Call the given callback on the main twisted thread after the transaction has
|
|
|
|
finished.
|
|
|
|
|
|
|
|
Mostly used to invalidate the caches on the correct thread.
|
|
|
|
|
|
|
|
Note that transactions may be retried a few times if they encounter database
|
|
|
|
errors such as serialization failures. Callbacks given to `call_after`
|
|
|
|
will accumulate across transaction attempts and will _all_ be called once a
|
|
|
|
transaction attempt succeeds, regardless of whether previous transaction
|
|
|
|
attempts failed. Otherwise, if all transaction attempts fail, all
|
|
|
|
`call_on_exception` callbacks will be run instead.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-02-27 06:53:40 -05:00
|
|
|
# if self.after_callbacks is None, that means that whatever constructed the
|
|
|
|
# LoggingTransaction isn't expecting there to be any callbacks; assert that
|
|
|
|
# is not the case.
|
|
|
|
assert self.after_callbacks is not None
|
2022-09-30 12:36:28 -04:00
|
|
|
self.after_callbacks.append((callback, args, kwargs))
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-07-19 07:25:29 -04:00
|
|
|
def async_call_after(
|
|
|
|
self, callback: Callable[P, Awaitable], *args: P.args, **kwargs: P.kwargs
|
|
|
|
) -> None:
|
|
|
|
"""Call the given asynchronous callback on the main twisted thread after
|
|
|
|
the transaction has finished (but before those added in `call_after`).
|
|
|
|
|
|
|
|
Mostly used to invalidate remote caches after transactions.
|
|
|
|
|
|
|
|
Note that transactions may be retried a few times if they encounter database
|
|
|
|
errors such as serialization failures. Callbacks given to `async_call_after`
|
|
|
|
will accumulate across transaction attempts and will _all_ be called once a
|
|
|
|
transaction attempt succeeds, regardless of whether previous transaction
|
|
|
|
attempts failed. Otherwise, if all transaction attempts fail, all
|
|
|
|
`call_on_exception` callbacks will be run instead.
|
|
|
|
"""
|
|
|
|
# if self.async_after_callbacks is None, that means that whatever constructed the
|
|
|
|
# LoggingTransaction isn't expecting there to be any callbacks; assert that
|
|
|
|
# is not the case.
|
|
|
|
assert self.async_after_callbacks is not None
|
2022-09-30 12:36:28 -04:00
|
|
|
self.async_after_callbacks.append((callback, args, kwargs))
|
2022-07-19 07:25:29 -04:00
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def call_on_exception(
|
2022-05-09 06:27:39 -04:00
|
|
|
self, callback: Callable[P, object], *args: P.args, **kwargs: P.kwargs
|
|
|
|
) -> None:
|
2022-03-29 07:31:05 -04:00
|
|
|
"""Call the given callback on the main twisted thread after the transaction has
|
|
|
|
failed.
|
|
|
|
|
|
|
|
Note that transactions may be retried a few times if they encounter database
|
|
|
|
errors such as serialization failures. Callbacks given to `call_on_exception`
|
|
|
|
will accumulate across transaction attempts and will _all_ be called once the
|
|
|
|
final transaction attempt fails. No `call_on_exception` callbacks will be run
|
|
|
|
if any transaction attempt succeeds.
|
|
|
|
"""
|
2020-02-27 06:53:40 -05:00
|
|
|
# if self.exception_callbacks is None, that means that whatever constructed the
|
|
|
|
# LoggingTransaction isn't expecting there to be any callbacks; assert that
|
|
|
|
# is not the case.
|
|
|
|
assert self.exception_callbacks is not None
|
2022-09-30 12:36:28 -04:00
|
|
|
self.exception_callbacks.append((callback, args, kwargs))
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2021-02-05 15:39:19 -05:00
|
|
|
def fetchone(self) -> Optional[Tuple]:
|
|
|
|
return self.txn.fetchone()
|
|
|
|
|
|
|
|
def fetchmany(self, size: Optional[int] = None) -> List[Tuple]:
|
|
|
|
return self.txn.fetchmany(size=size)
|
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
def fetchall(self) -> List[Tuple]:
|
|
|
|
return self.txn.fetchall()
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
def __iter__(self) -> Iterator[Tuple]:
|
2019-12-04 08:52:46 -05:00
|
|
|
return self.txn.__iter__()
|
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
@property
|
|
|
|
def rowcount(self) -> int:
|
|
|
|
return self.txn.rowcount
|
|
|
|
|
|
|
|
@property
|
2023-08-29 10:38:56 -04:00
|
|
|
def description(
|
|
|
|
self,
|
2023-09-25 12:48:42 -04:00
|
|
|
) -> Optional[Sequence[Any]]:
|
2020-02-27 06:53:40 -05:00
|
|
|
return self.txn.description
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def execute_batch(self, sql: str, args: Iterable[Iterable[Any]]) -> None:
|
2021-01-21 09:44:12 -05:00
|
|
|
"""Similar to `executemany`, except `txn.rowcount` will not be correct
|
|
|
|
afterwards.
|
|
|
|
|
|
|
|
More efficient than `executemany` on PostgreSQL
|
|
|
|
"""
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
if isinstance(self.database_engine, PostgresEngine):
|
2022-03-23 10:03:24 -04:00
|
|
|
from psycopg2.extras import execute_batch
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2023-04-14 14:04:49 -04:00
|
|
|
# TODO: is it safe for values to be Iterable[Iterable[Any]] here?
|
|
|
|
# https://www.psycopg.org/docs/extras.html?highlight=execute_batch#psycopg2.extras.execute_batch
|
|
|
|
# suggests each arg in args should be a sequence or mapping
|
2022-03-28 12:21:23 -04:00
|
|
|
self._do_execute(
|
|
|
|
lambda the_sql: execute_batch(self.txn, the_sql, args), sql
|
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
else:
|
2023-04-14 14:04:49 -04:00
|
|
|
# TODO: is it safe for values to be Iterable[Iterable[Any]] here?
|
|
|
|
# https://docs.python.org/3/library/sqlite3.html?highlight=sqlite3#sqlite3.Cursor.executemany
|
|
|
|
# suggests that the outer collection may be iterable, but
|
|
|
|
# https://docs.python.org/3/library/sqlite3.html?highlight=sqlite3#how-to-use-placeholders-to-bind-values-in-sql-queries
|
|
|
|
# suggests that the inner collection should be a sequence or dict.
|
2021-01-21 05:22:53 -05:00
|
|
|
self.executemany(sql, args)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-03-28 12:21:23 -04:00
|
|
|
def execute_values(
|
2023-05-03 09:41:37 -04:00
|
|
|
self,
|
|
|
|
sql: str,
|
|
|
|
values: Iterable[Iterable[Any]],
|
|
|
|
template: Optional[str] = None,
|
|
|
|
fetch: bool = True,
|
2022-03-28 12:21:23 -04:00
|
|
|
) -> List[Tuple]:
|
2021-01-11 11:09:22 -05:00
|
|
|
"""Corresponds to psycopg2.extras.execute_values. Only available when
|
|
|
|
using postgres.
|
|
|
|
|
2021-09-03 11:35:49 -04:00
|
|
|
The `fetch` parameter must be set to False if the query does not return
|
|
|
|
rows (e.g. INSERTs).
|
2023-05-03 09:41:37 -04:00
|
|
|
|
|
|
|
The `template` is the snippet to merge to every item in argslist to
|
|
|
|
compose the query.
|
2021-01-11 11:09:22 -05:00
|
|
|
"""
|
|
|
|
assert isinstance(self.database_engine, PostgresEngine)
|
2022-03-23 10:03:24 -04:00
|
|
|
from psycopg2.extras import execute_values
|
2021-01-11 11:09:22 -05:00
|
|
|
|
|
|
|
return self._do_execute(
|
2023-04-14 14:04:49 -04:00
|
|
|
# TODO: is it safe for values to be Iterable[Iterable[Any]] here?
|
|
|
|
# https://www.psycopg.org/docs/extras.html?highlight=execute_batch#psycopg2.extras.execute_values says values should be Sequence[Sequence]
|
2023-09-08 09:50:13 -04:00
|
|
|
lambda the_sql, the_values: execute_values(
|
|
|
|
self.txn, the_sql, the_values, template=template, fetch=fetch
|
2023-05-03 09:41:37 -04:00
|
|
|
),
|
2022-03-23 10:03:24 -04:00
|
|
|
sql,
|
2023-09-08 09:50:13 -04:00
|
|
|
values,
|
2021-01-11 11:09:22 -05:00
|
|
|
)
|
|
|
|
|
2023-04-14 14:04:49 -04:00
|
|
|
def execute(self, sql: str, parameters: SQLQueryParameters = ()) -> None:
|
|
|
|
self._do_execute(self.txn.execute, sql, parameters)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def executemany(self, sql: str, *args: Any) -> None:
|
2023-10-30 10:34:37 -04:00
|
|
|
"""Repeatedly execute the same piece of SQL with different parameters.
|
|
|
|
|
|
|
|
See https://peps.python.org/pep-0249/#executemany. Note in particular that
|
|
|
|
|
|
|
|
> Use of this method for an operation which produces one or more result sets
|
|
|
|
> constitutes undefined behavior
|
|
|
|
|
|
|
|
so you can't use this for e.g. a SELECT, an UPDATE ... RETURNING, or a
|
|
|
|
DELETE FROM... RETURNING.
|
|
|
|
"""
|
2023-04-14 14:04:49 -04:00
|
|
|
# TODO: we should add a type for *args here. Looking at Cursor.executemany
|
|
|
|
# and DBAPI2 it ought to be Sequence[_Parameter], but we pass in
|
|
|
|
# Iterable[Iterable[Any]] in execute_batch and execute_values above, which mypy
|
|
|
|
# complains about.
|
2019-12-04 08:52:46 -05:00
|
|
|
self._do_execute(self.txn.executemany, sql, *args)
|
|
|
|
|
2022-09-26 13:28:32 -04:00
|
|
|
def executescript(self, sql: str) -> None:
|
|
|
|
if isinstance(self.database_engine, Sqlite3Engine):
|
|
|
|
self._do_execute(self.txn.executescript, sql) # type: ignore[attr-defined]
|
|
|
|
else:
|
|
|
|
raise NotImplementedError(
|
|
|
|
f"executescript only exists for sqlite driver, not {type(self.database_engine)}"
|
|
|
|
)
|
|
|
|
|
2020-05-12 06:20:48 -04:00
|
|
|
def _make_sql_one_line(self, sql: str) -> str:
|
2019-12-04 08:52:46 -05:00
|
|
|
"Strip newlines out of SQL so that the loggers in the DB are on one line"
|
2020-05-12 06:20:48 -04:00
|
|
|
return " ".join(line.strip() for line in sql.splitlines() if line.strip())
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-05-09 05:28:38 -04:00
|
|
|
def _do_execute(
|
|
|
|
self,
|
|
|
|
func: Callable[Concatenate[str, P], R],
|
|
|
|
sql: str,
|
|
|
|
*args: P.args,
|
|
|
|
**kwargs: P.kwargs,
|
|
|
|
) -> R:
|
2022-06-30 04:43:24 -04:00
|
|
|
# Generate a one-line version of the SQL to better log it.
|
|
|
|
one_line_sql = self._make_sql_one_line(sql)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
# TODO(paul): Maybe use 'info' and 'debug' for values?
|
2022-06-30 04:43:24 -04:00
|
|
|
sql_logger.debug("[SQL] {%s} %s", self.name, one_line_sql)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
sql = self.database_engine.convert_param_style(sql)
|
|
|
|
if args:
|
|
|
|
try:
|
2022-09-30 12:36:28 -04:00
|
|
|
sql_logger.debug("[SQL values] {%s} %r", self.name, args[0])
|
2019-12-04 08:52:46 -05:00
|
|
|
except Exception:
|
|
|
|
# Don't let logging failures stop SQL from working
|
|
|
|
pass
|
|
|
|
|
|
|
|
start = time.time()
|
|
|
|
|
|
|
|
try:
|
2021-06-03 11:31:56 -04:00
|
|
|
with opentracing.start_active_span(
|
|
|
|
"db.query",
|
|
|
|
tags={
|
|
|
|
opentracing.tags.DATABASE_TYPE: "sql",
|
2022-06-30 04:43:24 -04:00
|
|
|
opentracing.tags.DATABASE_STATEMENT: one_line_sql,
|
2021-06-03 11:31:56 -04:00
|
|
|
},
|
|
|
|
):
|
2022-05-09 05:28:38 -04:00
|
|
|
return func(sql, *args, **kwargs)
|
2019-12-04 08:52:46 -05:00
|
|
|
except Exception as e:
|
2020-07-28 13:52:13 -04:00
|
|
|
sql_logger.debug("[SQL FAIL] {%s} %s", self.name, e)
|
2019-12-04 08:52:46 -05:00
|
|
|
raise
|
|
|
|
finally:
|
|
|
|
secs = time.time() - start
|
|
|
|
sql_logger.debug("[SQL time] {%s} %f sec", self.name, secs)
|
|
|
|
sql_query_timer.labels(sql.split()[0]).observe(secs)
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def close(self) -> None:
|
2020-02-27 06:53:40 -05:00
|
|
|
self.txn.close()
|
|
|
|
|
2020-10-02 10:20:45 -04:00
|
|
|
def __enter__(self) -> "LoggingTransaction":
|
|
|
|
return self
|
|
|
|
|
2022-05-12 10:33:50 -04:00
|
|
|
def __exit__(
|
|
|
|
self,
|
|
|
|
exc_type: Optional[Type[BaseException]],
|
|
|
|
exc_value: Optional[BaseException],
|
|
|
|
traceback: Optional[types.TracebackType],
|
|
|
|
) -> None:
|
2020-10-02 10:20:45 -04:00
|
|
|
self.close()
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-09-04 06:54:56 -04:00
|
|
|
class PerformanceCounters:
|
2022-05-12 10:33:50 -04:00
|
|
|
def __init__(self) -> None:
|
|
|
|
self.current_counters: Dict[str, Tuple[int, float]] = {}
|
|
|
|
self.previous_counters: Dict[str, Tuple[int, float]] = {}
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def update(self, key: str, duration_secs: float) -> None:
|
2022-05-12 10:33:50 -04:00
|
|
|
count, cum_time = self.current_counters.get(key, (0, 0.0))
|
2019-12-04 08:52:46 -05:00
|
|
|
count += 1
|
|
|
|
cum_time += duration_secs
|
|
|
|
self.current_counters[key] = (count, cum_time)
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def interval(self, interval_duration_secs: float, limit: int = 3) -> str:
|
2019-12-04 08:52:46 -05:00
|
|
|
counters = []
|
2020-06-15 07:03:36 -04:00
|
|
|
for name, (count, cum_time) in self.current_counters.items():
|
2019-12-04 08:52:46 -05:00
|
|
|
prev_count, prev_time = self.previous_counters.get(name, (0, 0))
|
|
|
|
counters.append(
|
|
|
|
(
|
|
|
|
(cum_time - prev_time) / interval_duration_secs,
|
|
|
|
count - prev_count,
|
|
|
|
name,
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
self.previous_counters = dict(self.current_counters)
|
|
|
|
|
|
|
|
counters.sort(reverse=True)
|
|
|
|
|
|
|
|
top_n_counters = ", ".join(
|
|
|
|
"%s(%d): %.3f%%" % (name, count, 100 * ratio)
|
|
|
|
for ratio, count, name in counters[:limit]
|
|
|
|
)
|
|
|
|
|
|
|
|
return top_n_counters
|
|
|
|
|
|
|
|
|
2020-09-04 06:54:56 -04:00
|
|
|
class DatabasePool:
|
2019-12-04 09:00:29 -05:00
|
|
|
"""Wraps a single physical database and connection pool.
|
|
|
|
|
|
|
|
A single database may be used by multiple data stores.
|
|
|
|
"""
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
_TXN_ID = 0
|
2023-02-09 09:49:04 -05:00
|
|
|
engine: BaseDatabaseEngine
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
def __init__(
|
2021-02-24 05:13:53 -05:00
|
|
|
self,
|
2021-10-22 13:15:41 -04:00
|
|
|
hs: "HomeServer",
|
2021-02-24 05:13:53 -05:00
|
|
|
database_config: DatabaseConnectionConfig,
|
|
|
|
engine: BaseDatabaseEngine,
|
2020-02-27 06:53:40 -05:00
|
|
|
):
|
2019-12-04 08:52:46 -05:00
|
|
|
self.hs = hs
|
|
|
|
self._clock = hs.get_clock()
|
2021-08-02 09:24:43 -04:00
|
|
|
self._txn_limit = database_config.config.get("txn_limit", 0)
|
2019-12-18 05:45:12 -05:00
|
|
|
self._database_config = database_config
|
|
|
|
self._db_pool = make_pool(hs.get_reactor(), database_config, engine)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2019-12-04 10:09:36 -05:00
|
|
|
self.updates = BackgroundUpdater(hs, self)
|
2023-06-07 13:12:23 -04:00
|
|
|
LaterGauge(
|
|
|
|
"synapse_background_update_status",
|
|
|
|
"Background update status",
|
|
|
|
[],
|
|
|
|
self.updates.get_status,
|
|
|
|
)
|
2019-12-04 10:09:36 -05:00
|
|
|
|
2020-02-27 06:53:40 -05:00
|
|
|
self._previous_txn_total_time = 0.0
|
|
|
|
self._current_txn_total_time = 0.0
|
|
|
|
self._previous_loop_ts = 0.0
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2021-08-02 09:24:43 -04:00
|
|
|
# Transaction counter: key is the twisted thread id, value is the current count
|
|
|
|
self._txn_counters: Dict[int, int] = defaultdict(int)
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
# TODO(paul): These can eventually be removed once the metrics code
|
|
|
|
# is running in mainline, and we have some nice monitoring frontends
|
|
|
|
# to watch it
|
|
|
|
self._txn_perf_counters = PerformanceCounters()
|
|
|
|
|
2019-12-18 05:45:12 -05:00
|
|
|
self.engine = engine
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
# A set of tables that are not safe to use native upserts in.
|
|
|
|
self._unsafe_to_upsert_tables = set(UNIQUE_INDEX_BACKGROUND_UPDATES.keys())
|
|
|
|
|
2023-05-19 08:25:25 -04:00
|
|
|
# The user_directory_search table is unsafe to use native upserts
|
|
|
|
# on SQLite because the existing search table does not have an index.
|
2019-12-06 08:40:02 -05:00
|
|
|
if isinstance(self.engine, Sqlite3Engine):
|
2019-12-04 08:52:46 -05:00
|
|
|
self._unsafe_to_upsert_tables.add("user_directory_search")
|
|
|
|
|
2022-09-09 06:14:10 -04:00
|
|
|
# Check ASAP (and then later, every 1s) to see if we have finished
|
|
|
|
# background updates of tables that aren't safe to update.
|
|
|
|
self._clock.call_later(
|
|
|
|
0.0,
|
|
|
|
run_as_background_process,
|
|
|
|
"upsert_safety_check",
|
|
|
|
self._check_safe_to_upsert,
|
|
|
|
)
|
2021-01-14 12:19:35 -05:00
|
|
|
|
2021-11-08 11:08:02 -05:00
|
|
|
def name(self) -> str:
|
|
|
|
"Return the name of this database"
|
|
|
|
return self._database_config.name
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def is_running(self) -> bool:
|
2021-02-16 17:32:34 -05:00
|
|
|
"""Is the database pool currently running"""
|
2019-12-18 05:45:12 -05:00
|
|
|
return self._db_pool.running
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
async def _check_safe_to_upsert(self) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Is it safe to use native UPSERT?
|
|
|
|
|
|
|
|
If there are background updates, we will need to wait, as they may be
|
|
|
|
the addition of indexes that set the UNIQUE constraint that we require.
|
|
|
|
|
|
|
|
If the background updates have not completed, wait 15 sec and check again.
|
|
|
|
"""
|
2023-10-26 13:01:36 -04:00
|
|
|
updates = cast(
|
|
|
|
List[Tuple[str]],
|
|
|
|
await self.simple_select_list(
|
|
|
|
"background_updates",
|
|
|
|
keyvalues=None,
|
|
|
|
retcols=["update_name"],
|
|
|
|
desc="check_background_updates",
|
|
|
|
),
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
2023-10-26 13:01:36 -04:00
|
|
|
background_update_names = [x[0] for x in updates]
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
for table, update_name in UNIQUE_INDEX_BACKGROUND_UPDATES.items():
|
2022-11-17 11:09:56 -05:00
|
|
|
if update_name not in background_update_names:
|
2019-12-04 08:52:46 -05:00
|
|
|
logger.debug("Now safe to upsert in %s", table)
|
|
|
|
self._unsafe_to_upsert_tables.discard(table)
|
|
|
|
|
|
|
|
# If there's any updates still running, reschedule to run.
|
2022-11-17 11:09:56 -05:00
|
|
|
if background_update_names:
|
2019-12-04 08:52:46 -05:00
|
|
|
self._clock.call_later(
|
|
|
|
15.0,
|
|
|
|
run_as_background_process,
|
|
|
|
"upsert_safety_check",
|
|
|
|
self._check_safe_to_upsert,
|
|
|
|
)
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
def start_profiling(self) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
self._previous_loop_ts = monotonic_time()
|
|
|
|
|
2022-05-12 10:33:50 -04:00
|
|
|
def loop() -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
curr = self._current_txn_total_time
|
|
|
|
prev = self._previous_txn_total_time
|
|
|
|
self._previous_txn_total_time = curr
|
|
|
|
|
|
|
|
time_now = monotonic_time()
|
|
|
|
time_then = self._previous_loop_ts
|
|
|
|
self._previous_loop_ts = time_now
|
|
|
|
|
|
|
|
duration = time_now - time_then
|
|
|
|
ratio = (curr - prev) / duration
|
|
|
|
|
|
|
|
top_three_counters = self._txn_perf_counters.interval(duration, limit=3)
|
|
|
|
|
2020-02-05 03:57:38 -05:00
|
|
|
perf_logger.debug(
|
2019-12-04 08:52:46 -05:00
|
|
|
"Total database time: %.3f%% {%s}", ratio * 100, top_three_counters
|
|
|
|
)
|
|
|
|
|
|
|
|
self._clock.looping_call(loop, 10000)
|
|
|
|
|
|
|
|
def new_transaction(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
2020-10-02 10:20:45 -04:00
|
|
|
conn: LoggingDatabaseConnection,
|
2020-08-20 09:00:59 -04:00
|
|
|
desc: str,
|
|
|
|
after_callbacks: List[_CallbackListEntry],
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks: List[_AsyncCallbackListEntry],
|
2020-08-20 09:00:59 -04:00
|
|
|
exception_callbacks: List[_CallbackListEntry],
|
2022-05-09 06:27:39 -04:00
|
|
|
func: Callable[Concatenate[LoggingTransaction, P], R],
|
|
|
|
*args: P.args,
|
|
|
|
**kwargs: P.kwargs,
|
2020-08-20 09:00:59 -04:00
|
|
|
) -> R:
|
2020-10-07 10:15:57 -04:00
|
|
|
"""Start a new database transaction with the given connection.
|
|
|
|
|
|
|
|
Note: The given func may be called multiple times under certain
|
|
|
|
failure modes. This is normally fine when in a standard transaction,
|
|
|
|
but care must be taken if the connection is in `autocommit` mode that
|
|
|
|
the function will correctly handle being aborted and retried half way
|
|
|
|
through its execution.
|
|
|
|
|
2021-12-13 14:01:27 -05:00
|
|
|
Similarly, the arguments to `func` (`args`, `kwargs`) should not be generators,
|
|
|
|
since they could be evaluated multiple times (which would produce an empty
|
|
|
|
result on the second or subsequent evaluation). Likewise, the closure of `func`
|
|
|
|
must not reference any generators. This method attempts to detect such usage
|
|
|
|
and will log an error.
|
|
|
|
|
2020-10-07 10:15:57 -04:00
|
|
|
Args:
|
|
|
|
conn
|
|
|
|
desc
|
|
|
|
after_callbacks
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks
|
2020-10-07 10:15:57 -04:00
|
|
|
exception_callbacks
|
|
|
|
func
|
|
|
|
*args
|
|
|
|
**kwargs
|
|
|
|
"""
|
|
|
|
|
2021-12-13 14:01:27 -05:00
|
|
|
# Robustness check: ensure that none of the arguments are generators, since that
|
|
|
|
# will fail if we have to repeat the transaction.
|
|
|
|
# For now, we just log an error, and hope that it works on the first attempt.
|
|
|
|
# TODO: raise an exception.
|
2022-05-09 06:27:39 -04:00
|
|
|
|
2022-09-30 12:36:28 -04:00
|
|
|
for i, arg in enumerate(args):
|
2021-12-13 14:01:27 -05:00
|
|
|
if inspect.isgenerator(arg):
|
|
|
|
logger.error(
|
|
|
|
"Programming error: generator passed to new_transaction as "
|
|
|
|
"argument %i to function %s",
|
|
|
|
i,
|
|
|
|
func,
|
|
|
|
)
|
2022-09-30 12:36:28 -04:00
|
|
|
for name, val in kwargs.items():
|
2021-12-13 14:01:27 -05:00
|
|
|
if inspect.isgenerator(val):
|
|
|
|
logger.error(
|
|
|
|
"Programming error: generator passed to new_transaction as "
|
|
|
|
"argument %s to function %s",
|
|
|
|
name,
|
|
|
|
func,
|
|
|
|
)
|
|
|
|
# also check variables referenced in func's closure
|
|
|
|
if inspect.isfunction(func):
|
2022-12-12 16:25:07 -05:00
|
|
|
# Keep the cast for now---it helps PyCharm to understand what `func` is.
|
|
|
|
f = cast(types.FunctionType, func) # type: ignore[redundant-cast]
|
2021-12-13 14:01:27 -05:00
|
|
|
if f.__closure__:
|
|
|
|
for i, cell in enumerate(f.__closure__):
|
2023-03-07 03:51:34 -05:00
|
|
|
try:
|
|
|
|
contents = cell.cell_contents
|
|
|
|
except ValueError:
|
|
|
|
# cell.cell_contents can raise if the "cell" is empty,
|
|
|
|
# which indicates that the variable is currently
|
|
|
|
# unbound.
|
|
|
|
continue
|
|
|
|
|
|
|
|
if inspect.isgenerator(contents):
|
2021-12-13 14:01:27 -05:00
|
|
|
logger.error(
|
|
|
|
"Programming error: function %s references generator %s "
|
|
|
|
"via its closure",
|
|
|
|
f,
|
|
|
|
f.__code__.co_freevars[i],
|
|
|
|
)
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
start = monotonic_time()
|
|
|
|
txn_id = self._TXN_ID
|
|
|
|
|
|
|
|
# We don't really need these to be unique, so lets stop it from
|
|
|
|
# growing really large.
|
|
|
|
self._TXN_ID = (self._TXN_ID + 1) % (MAX_TXN_ID)
|
|
|
|
|
|
|
|
name = "%s-%x" % (desc, txn_id)
|
|
|
|
|
|
|
|
transaction_logger.debug("[TXN START] {%s}", name)
|
|
|
|
|
|
|
|
try:
|
|
|
|
i = 0
|
|
|
|
N = 5
|
|
|
|
while True:
|
2020-10-02 10:20:45 -04:00
|
|
|
cursor = conn.cursor(
|
|
|
|
txn_name=name,
|
|
|
|
after_callbacks=after_callbacks,
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks=async_after_callbacks,
|
2020-10-02 10:20:45 -04:00
|
|
|
exception_callbacks=exception_callbacks,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
try:
|
2021-06-03 11:31:56 -04:00
|
|
|
with opentracing.start_active_span(
|
|
|
|
"db.txn",
|
|
|
|
tags={
|
|
|
|
opentracing.SynapseTags.DB_TXN_DESC: desc,
|
|
|
|
opentracing.SynapseTags.DB_TXN_ID: name,
|
|
|
|
},
|
|
|
|
):
|
|
|
|
r = func(cursor, *args, **kwargs)
|
2021-06-07 13:01:32 -04:00
|
|
|
opentracing.log_kv({"message": "commit"})
|
2021-06-03 11:31:56 -04:00
|
|
|
conn.commit()
|
|
|
|
return r
|
2019-12-06 08:40:02 -05:00
|
|
|
except self.engine.module.OperationalError as e:
|
2019-12-04 08:52:46 -05:00
|
|
|
# This can happen if the database disappears mid
|
|
|
|
# transaction.
|
2020-07-28 13:52:13 -04:00
|
|
|
transaction_logger.warning(
|
2021-02-16 17:32:34 -05:00
|
|
|
"[TXN OPERROR] {%s} %s %d/%d",
|
|
|
|
name,
|
|
|
|
e,
|
|
|
|
i,
|
|
|
|
N,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
if i < N:
|
|
|
|
i += 1
|
|
|
|
try:
|
2021-06-07 13:01:32 -04:00
|
|
|
with opentracing.start_active_span("db.rollback"):
|
|
|
|
conn.rollback()
|
2019-12-06 08:40:02 -05:00
|
|
|
except self.engine.module.Error as e1:
|
2020-07-28 13:52:13 -04:00
|
|
|
transaction_logger.warning("[TXN EROLL] {%s} %s", name, e1)
|
2019-12-04 08:52:46 -05:00
|
|
|
continue
|
|
|
|
raise
|
2019-12-06 08:40:02 -05:00
|
|
|
except self.engine.module.DatabaseError as e:
|
|
|
|
if self.engine.is_deadlock(e):
|
2020-07-28 13:52:13 -04:00
|
|
|
transaction_logger.warning(
|
|
|
|
"[TXN DEADLOCK] {%s} %d/%d", name, i, N
|
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
if i < N:
|
|
|
|
i += 1
|
|
|
|
try:
|
2021-06-07 13:01:32 -04:00
|
|
|
with opentracing.start_active_span("db.rollback"):
|
|
|
|
conn.rollback()
|
2019-12-06 08:40:02 -05:00
|
|
|
except self.engine.module.Error as e1:
|
2020-07-28 13:52:13 -04:00
|
|
|
transaction_logger.warning(
|
2021-02-16 17:32:34 -05:00
|
|
|
"[TXN EROLL] {%s} %s",
|
|
|
|
name,
|
|
|
|
e1,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
continue
|
|
|
|
raise
|
|
|
|
finally:
|
|
|
|
# we're either about to retry with a new cursor, or we're about to
|
|
|
|
# release the connection. Once we release the connection, it could
|
|
|
|
# get used for another query, which might do a conn.rollback().
|
|
|
|
#
|
|
|
|
# In the latter case, even though that probably wouldn't affect the
|
|
|
|
# results of this transaction, python's sqlite will reset all
|
|
|
|
# statements on the connection [1], which will make our cursor
|
|
|
|
# invalid [2].
|
|
|
|
#
|
|
|
|
# In any case, continuing to read rows after commit()ing seems
|
|
|
|
# dubious from the PoV of ACID transactional semantics
|
|
|
|
# (sqlite explicitly says that once you commit, you may see rows
|
|
|
|
# from subsequent updates.)
|
|
|
|
#
|
|
|
|
# In psycopg2, cursors are essentially a client-side fabrication -
|
|
|
|
# all the data is transferred to the client side when the statement
|
|
|
|
# finishes executing - so in theory we could go on streaming results
|
|
|
|
# from the cursor, but attempting to do so would make us
|
|
|
|
# incompatible with sqlite, so let's make sure we're not doing that
|
|
|
|
# by closing the cursor.
|
|
|
|
#
|
|
|
|
# (*named* cursors in psycopg2 are different and are proper server-
|
|
|
|
# side things, but (a) we don't use them and (b) they are implicitly
|
|
|
|
# closed by ending the transaction anyway.)
|
|
|
|
#
|
|
|
|
# In short, if we haven't finished with the cursor yet, that's a
|
|
|
|
# problem waiting to bite us.
|
|
|
|
#
|
|
|
|
# TL;DR: we're done with the cursor, so we can close it.
|
|
|
|
#
|
|
|
|
# [1]: https://github.com/python/cpython/blob/v3.8.0/Modules/_sqlite/connection.c#L465
|
|
|
|
# [2]: https://github.com/python/cpython/blob/v3.8.0/Modules/_sqlite/cursor.c#L236
|
|
|
|
cursor.close()
|
|
|
|
except Exception as e:
|
2020-07-28 13:52:13 -04:00
|
|
|
transaction_logger.debug("[TXN FAIL] {%s} %s", name, e)
|
2019-12-04 08:52:46 -05:00
|
|
|
raise
|
|
|
|
finally:
|
|
|
|
end = monotonic_time()
|
|
|
|
duration = end - start
|
|
|
|
|
2020-03-24 10:45:33 -04:00
|
|
|
current_context().add_database_transaction(duration)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
transaction_logger.debug("[TXN END] {%s} %f sec", name, duration)
|
|
|
|
|
|
|
|
self._current_txn_total_time += duration
|
|
|
|
self._txn_perf_counters.update(desc, duration)
|
2022-07-21 08:23:05 -04:00
|
|
|
sql_txn_count.labels(desc).inc(1)
|
|
|
|
sql_txn_duration.labels(desc).inc(duration)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-09-02 13:11:02 -04:00
|
|
|
async def runInteraction(
|
2020-10-07 10:15:57 -04:00
|
|
|
self,
|
|
|
|
desc: str,
|
2021-04-27 08:13:07 -04:00
|
|
|
func: Callable[..., R],
|
2020-10-07 10:15:57 -04:00
|
|
|
*args: Any,
|
|
|
|
db_autocommit: bool = False,
|
2022-01-25 09:14:46 -05:00
|
|
|
isolation_level: Optional[int] = None,
|
2021-04-13 05:41:34 -04:00
|
|
|
**kwargs: Any,
|
2020-09-02 13:11:02 -04:00
|
|
|
) -> R:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Starts a transaction on the database and runs a given function
|
|
|
|
|
|
|
|
Arguments:
|
2020-02-27 06:53:40 -05:00
|
|
|
desc: description of the transaction, for logging and metrics
|
|
|
|
func: callback function, which will be called with a
|
2019-12-04 08:52:46 -05:00
|
|
|
database transaction (twisted.enterprise.adbapi.Transaction) as
|
|
|
|
its first argument, followed by `args` and `kwargs`.
|
|
|
|
|
2020-10-07 10:15:57 -04:00
|
|
|
db_autocommit: Whether to run the function in "autocommit" mode,
|
|
|
|
i.e. outside of a transaction. This is useful for transactions
|
|
|
|
that are only a single query.
|
|
|
|
|
|
|
|
Currently, this is only implemented for Postgres. SQLite will still
|
|
|
|
run the function inside a transaction.
|
|
|
|
|
|
|
|
WARNING: This means that if func fails half way through then
|
|
|
|
the changes will *not* be rolled back. `func` may also get
|
|
|
|
called multiple times if the transaction is retried, so must
|
|
|
|
correctly handle that case.
|
|
|
|
|
2022-01-25 09:14:46 -05:00
|
|
|
isolation_level: Set the server isolation level for this transaction.
|
2020-02-27 06:53:40 -05:00
|
|
|
args: positional args to pass to `func`
|
|
|
|
kwargs: named args to pass to `func`
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Returns:
|
2020-09-02 13:11:02 -04:00
|
|
|
The result of func
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
|
2022-03-16 11:07:41 -04:00
|
|
|
async def _runInteraction() -> R:
|
|
|
|
after_callbacks: List[_CallbackListEntry] = []
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks: List[_AsyncCallbackListEntry] = []
|
2022-03-16 11:07:41 -04:00
|
|
|
exception_callbacks: List[_CallbackListEntry] = []
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-03-16 11:07:41 -04:00
|
|
|
if not current_context():
|
|
|
|
logger.warning("Starting db txn '%s' from sentinel context", desc)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-03-16 11:07:41 -04:00
|
|
|
try:
|
|
|
|
with opentracing.start_active_span(f"db.{desc}"):
|
2024-03-13 13:05:57 -04:00
|
|
|
result: R = await self.runWithConnection(
|
2023-03-09 07:10:09 -05:00
|
|
|
# mypy seems to have an issue with this, maybe a bug?
|
2024-03-13 13:05:57 -04:00
|
|
|
self.new_transaction,
|
2022-03-16 11:07:41 -04:00
|
|
|
desc,
|
|
|
|
after_callbacks,
|
2022-07-19 07:25:29 -04:00
|
|
|
async_after_callbacks,
|
2022-03-16 11:07:41 -04:00
|
|
|
exception_callbacks,
|
|
|
|
func,
|
|
|
|
*args,
|
|
|
|
db_autocommit=db_autocommit,
|
|
|
|
isolation_level=isolation_level,
|
|
|
|
**kwargs,
|
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-07-19 07:25:29 -04:00
|
|
|
# We order these assuming that async functions call out to external
|
|
|
|
# systems (e.g. to invalidate a cache) and the sync functions make these
|
|
|
|
# changes on any local in-memory caches/similar, and thus must be second.
|
|
|
|
for async_callback, async_args, async_kwargs in async_after_callbacks:
|
|
|
|
await async_callback(*async_args, **async_kwargs)
|
2022-03-16 11:07:41 -04:00
|
|
|
for after_callback, after_args, after_kwargs in after_callbacks:
|
2022-07-19 07:25:29 -04:00
|
|
|
after_callback(*after_args, **after_kwargs)
|
2024-03-13 13:05:57 -04:00
|
|
|
return result
|
2022-03-16 11:07:41 -04:00
|
|
|
except Exception:
|
2022-07-15 05:30:46 -04:00
|
|
|
for exception_callback, after_args, after_kwargs in exception_callbacks:
|
2022-07-19 07:25:29 -04:00
|
|
|
exception_callback(*after_args, **after_kwargs)
|
2022-03-16 11:07:41 -04:00
|
|
|
raise
|
|
|
|
|
|
|
|
# To handle cancellation, we ensure that `after_callback`s and
|
|
|
|
# `exception_callback`s are always run, since the transaction will complete
|
|
|
|
# on another thread regardless of cancellation.
|
|
|
|
#
|
|
|
|
# We also wait until everything above is done before releasing the
|
|
|
|
# `CancelledError`, so that logging contexts won't get used after they have been
|
|
|
|
# finished.
|
2022-04-22 13:20:06 -04:00
|
|
|
return await delay_cancellation(_runInteraction())
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
async def runWithConnection(
|
2020-10-07 10:15:57 -04:00
|
|
|
self,
|
2023-03-09 07:10:09 -05:00
|
|
|
func: Callable[Concatenate[LoggingDatabaseConnection, P], R],
|
2020-10-07 10:15:57 -04:00
|
|
|
*args: Any,
|
|
|
|
db_autocommit: bool = False,
|
2022-01-25 09:14:46 -05:00
|
|
|
isolation_level: Optional[int] = None,
|
2021-04-13 05:41:34 -04:00
|
|
|
**kwargs: Any,
|
2020-08-20 09:00:59 -04:00
|
|
|
) -> R:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Wraps the .runWithConnection() method on the underlying db_pool.
|
|
|
|
|
|
|
|
Arguments:
|
2020-02-27 06:53:40 -05:00
|
|
|
func: callback function, which will be called with a
|
2019-12-04 08:52:46 -05:00
|
|
|
database connection (twisted.enterprise.adbapi.Connection) as
|
|
|
|
its first argument, followed by `args` and `kwargs`.
|
2020-02-27 06:53:40 -05:00
|
|
|
args: positional args to pass to `func`
|
2020-10-07 10:15:57 -04:00
|
|
|
db_autocommit: Whether to run the function in "autocommit" mode,
|
|
|
|
i.e. outside of a transaction. This is useful for transaction
|
|
|
|
that are only a single query. Currently only affects postgres.
|
2022-01-25 09:14:46 -05:00
|
|
|
isolation_level: Set the server isolation level for this transaction.
|
2020-02-27 06:53:40 -05:00
|
|
|
kwargs: named args to pass to `func`
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Returns:
|
2020-08-19 07:09:24 -04:00
|
|
|
The result of func
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-01-05 08:06:55 -05:00
|
|
|
curr_context = current_context()
|
|
|
|
if not curr_context:
|
2019-12-04 08:52:46 -05:00
|
|
|
logger.warning(
|
|
|
|
"Starting db connection from sentinel context: metrics will be lost"
|
|
|
|
)
|
|
|
|
parent_context = None
|
2021-01-05 08:06:55 -05:00
|
|
|
else:
|
|
|
|
assert isinstance(curr_context, LoggingContext)
|
|
|
|
parent_context = curr_context
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
start_time = monotonic_time()
|
|
|
|
|
2023-03-09 07:10:09 -05:00
|
|
|
def inner_func(conn: _PoolConnection, *args: P.args, **kwargs: P.kwargs) -> R:
|
2020-10-07 10:15:57 -04:00
|
|
|
# We shouldn't be in a transaction. If we are then something
|
|
|
|
# somewhere hasn't committed after doing work. (This is likely only
|
|
|
|
# possible during startup, as `run*` will ensure changes are
|
|
|
|
# committed/rolled back before putting the connection back in the
|
|
|
|
# pool).
|
|
|
|
assert not self.engine.in_transaction(conn)
|
|
|
|
|
2021-04-28 07:44:52 -04:00
|
|
|
with LoggingContext(
|
|
|
|
str(curr_context), parent_context=parent_context
|
|
|
|
) as context:
|
2021-06-03 11:31:56 -04:00
|
|
|
with opentracing.start_active_span(
|
|
|
|
operation_name="db.connection",
|
|
|
|
):
|
|
|
|
sched_duration_sec = monotonic_time() - start_time
|
|
|
|
sql_scheduling_timer.observe(sched_duration_sec)
|
|
|
|
context.add_database_scheduled(sched_duration_sec)
|
|
|
|
|
2021-08-02 09:24:43 -04:00
|
|
|
if self._txn_limit > 0:
|
|
|
|
tid = self._db_pool.threadID()
|
|
|
|
self._txn_counters[tid] += 1
|
|
|
|
|
|
|
|
if self._txn_counters[tid] > self._txn_limit:
|
|
|
|
logger.debug(
|
|
|
|
"Reconnecting database connection over transaction limit"
|
|
|
|
)
|
|
|
|
conn.reconnect()
|
|
|
|
opentracing.log_kv(
|
|
|
|
{"message": "reconnected due to txn limit"}
|
|
|
|
)
|
|
|
|
self._txn_counters[tid] = 1
|
|
|
|
|
2021-06-03 11:31:56 -04:00
|
|
|
if self.engine.is_connection_closed(conn):
|
|
|
|
logger.debug("Reconnecting closed database connection")
|
|
|
|
conn.reconnect()
|
|
|
|
opentracing.log_kv({"message": "reconnected"})
|
2021-08-02 09:24:43 -04:00
|
|
|
if self._txn_limit > 0:
|
|
|
|
self._txn_counters[tid] = 1
|
2021-06-03 11:31:56 -04:00
|
|
|
|
|
|
|
try:
|
|
|
|
if db_autocommit:
|
|
|
|
self.engine.attempt_to_set_autocommit(conn, True)
|
2022-01-25 09:14:46 -05:00
|
|
|
if isolation_level is not None:
|
|
|
|
self.engine.attempt_to_set_isolation_level(
|
|
|
|
conn, isolation_level
|
|
|
|
)
|
2021-06-03 11:31:56 -04:00
|
|
|
|
|
|
|
db_conn = LoggingDatabaseConnection(
|
|
|
|
conn, self.engine, "runWithConnection"
|
|
|
|
)
|
|
|
|
return func(db_conn, *args, **kwargs)
|
|
|
|
finally:
|
|
|
|
if db_autocommit:
|
|
|
|
self.engine.attempt_to_set_autocommit(conn, False)
|
2022-01-25 09:14:46 -05:00
|
|
|
if isolation_level:
|
|
|
|
self.engine.attempt_to_set_isolation_level(conn, None)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-19 07:09:24 -04:00
|
|
|
return await make_deferred_yieldable(
|
2019-12-04 08:52:46 -05:00
|
|
|
self._db_pool.runWithConnection(inner_func, *args, **kwargs)
|
|
|
|
)
|
|
|
|
|
2023-10-26 15:12:28 -04:00
|
|
|
async def execute(self, desc: str, query: str, *args: Any) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Runs a single query for a result set.
|
|
|
|
|
|
|
|
Args:
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
query - The query string to execute
|
|
|
|
*args - Query args.
|
|
|
|
Returns:
|
2020-08-27 07:41:01 -04:00
|
|
|
The result of decoder(results)
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
|
2023-10-26 15:12:28 -04:00
|
|
|
def interaction(txn: LoggingTransaction) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
txn.execute(query, args)
|
2023-10-26 15:12:28 -04:00
|
|
|
return txn.fetchall()
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-27 07:41:01 -04:00
|
|
|
return await self.runInteraction(desc, interaction)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
# "Simple" SQL API methods that operate on a single table with no JOINs,
|
|
|
|
# no complex WHERE clauses, just a dict of values for columns.
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
async def simple_insert(
|
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
values: Dict[str, Any],
|
|
|
|
desc: str = "simple_insert",
|
2021-07-22 07:39:50 -04:00
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes an INSERT query on the named table.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
values: dict of new column names and values for them
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-07-22 07:39:50 -04:00
|
|
|
await self.runInteraction(desc, self.simple_insert_txn, table, values)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_insert_txn(
|
|
|
|
txn: LoggingTransaction, table: str, values: Dict[str, Any]
|
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
keys, vals = zip(*values.items())
|
|
|
|
|
|
|
|
sql = "INSERT INTO %s (%s) VALUES(%s)" % (
|
|
|
|
table,
|
|
|
|
", ".join(k for k in keys),
|
|
|
|
", ".join("?" for _ in keys),
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute(sql, vals)
|
|
|
|
|
2020-08-27 07:41:01 -04:00
|
|
|
async def simple_insert_many(
|
2021-12-10 10:02:33 -05:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keys: Collection[str],
|
2021-12-15 12:00:50 -05:00
|
|
|
values: Collection[Collection[Any]],
|
2021-12-10 10:02:33 -05:00
|
|
|
desc: str,
|
|
|
|
) -> None:
|
|
|
|
"""Executes an INSERT query on the named table.
|
|
|
|
|
|
|
|
The input is given as a list of rows, where each row is a list of values.
|
|
|
|
(Actually any iterable is fine.)
|
|
|
|
|
|
|
|
Args:
|
|
|
|
table: string giving the table name
|
|
|
|
keys: list of column names
|
|
|
|
values: for each row, a list of values in the same order as `keys`
|
|
|
|
desc: description of the transaction, for logging and metrics
|
|
|
|
"""
|
|
|
|
await self.runInteraction(
|
2022-01-13 19:44:18 -05:00
|
|
|
desc, self.simple_insert_many_txn, table, keys, values
|
2021-12-10 10:02:33 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
@staticmethod
|
2022-01-13 19:44:18 -05:00
|
|
|
def simple_insert_many_txn(
|
2021-12-10 10:02:33 -05:00
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
2023-11-09 10:57:09 -05:00
|
|
|
keys: Sequence[str],
|
2023-11-07 14:00:25 -05:00
|
|
|
values: Collection[Iterable[Any]],
|
2021-12-10 10:02:33 -05:00
|
|
|
) -> None:
|
|
|
|
"""Executes an INSERT query on the named table.
|
|
|
|
|
|
|
|
The input is given as a list of rows, where each row is a list of values.
|
|
|
|
(Actually any iterable is fine.)
|
|
|
|
|
|
|
|
Args:
|
|
|
|
txn: The transaction to use.
|
|
|
|
table: string giving the table name
|
|
|
|
keys: list of column names
|
|
|
|
values: for each row, a list of values in the same order as `keys`
|
|
|
|
"""
|
2023-11-07 14:00:25 -05:00
|
|
|
# If there's nothing to insert, then skip executing the query.
|
|
|
|
if not values:
|
|
|
|
return
|
2021-12-10 10:02:33 -05:00
|
|
|
|
2021-09-03 11:35:49 -04:00
|
|
|
if isinstance(txn.database_engine, PostgresEngine):
|
|
|
|
# We use `execute_values` as it can be a lot faster than `execute_batch`,
|
|
|
|
# but it's only available on postgres.
|
|
|
|
sql = "INSERT INTO %s (%s) VALUES ?" % (
|
|
|
|
table,
|
2021-12-10 10:02:33 -05:00
|
|
|
", ".join(k for k in keys),
|
2021-09-03 11:35:49 -04:00
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2021-12-10 10:02:33 -05:00
|
|
|
txn.execute_values(sql, values, fetch=False)
|
2021-09-03 11:35:49 -04:00
|
|
|
else:
|
|
|
|
sql = "INSERT INTO %s (%s) VALUES(%s)" % (
|
|
|
|
table,
|
2021-12-10 10:02:33 -05:00
|
|
|
", ".join(k for k in keys),
|
|
|
|
", ".join("?" for _ in keys),
|
2021-09-03 11:35:49 -04:00
|
|
|
)
|
|
|
|
|
2021-12-10 10:02:33 -05:00
|
|
|
txn.execute_batch(sql, values)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-17 12:18:01 -04:00
|
|
|
async def simple_upsert(
|
2019-12-04 08:52:46 -05:00
|
|
|
self,
|
2020-08-20 09:00:59 -04:00
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
values: Dict[str, Any],
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values: Optional[Dict[str, Any]] = None,
|
2023-09-15 04:16:45 -04:00
|
|
|
where_clause: Optional[str] = None,
|
2020-08-20 09:00:59 -04:00
|
|
|
desc: str = "simple_upsert",
|
2021-07-22 07:39:50 -04:00
|
|
|
) -> bool:
|
2022-09-29 14:10:47 -04:00
|
|
|
"""Insert a row with values + insertion_values; on conflict, update with values.
|
|
|
|
|
|
|
|
All of our supported databases accept the nonstandard "upsert" statement in
|
|
|
|
their dialect of SQL. We call this a "native upsert". The syntax looks roughly
|
|
|
|
like:
|
|
|
|
|
|
|
|
INSERT INTO table VALUES (values + insertion_values)
|
|
|
|
ON CONFLICT (keyvalues)
|
|
|
|
DO UPDATE SET (values); -- overwrite `values` columns only
|
|
|
|
|
|
|
|
If (values) is empty, the resulting query is slighlty simpler:
|
|
|
|
|
|
|
|
INSERT INTO table VALUES (insertion_values)
|
|
|
|
ON CONFLICT (keyvalues)
|
|
|
|
DO NOTHING; -- do not overwrite any columns
|
|
|
|
|
|
|
|
This function is a helper to build such queries.
|
|
|
|
|
|
|
|
In order for upserts to make sense, the database must be able to determine when
|
|
|
|
an upsert CONFLICTs with an existing row. Postgres and SQLite ensure this by
|
|
|
|
requiring that a unique index exist on the column names used to detect a
|
|
|
|
conflict (i.e. `keyvalues.keys()`).
|
|
|
|
|
2022-11-28 08:42:06 -05:00
|
|
|
If there is no such index yet[*], we can "emulate" an upsert with a SELECT
|
|
|
|
followed by either an INSERT or an UPDATE. This is unsafe unless *all* upserters
|
|
|
|
run at the SERIALIZABLE isolation level: we cannot make the same atomicity
|
|
|
|
guarantees that a native upsert can and are very vulnerable to races and
|
|
|
|
crashes. Therefore to upsert without an appropriate unique index, we acquire a
|
|
|
|
table-level lock before the emulated upsert.
|
2022-09-29 14:10:47 -04:00
|
|
|
|
|
|
|
[*]: Some tables have unique indices added to them in the background. Those
|
|
|
|
tables `T` are keys in the dictionary UNIQUE_INDEX_BACKGROUND_UPDATES,
|
|
|
|
where `T` maps to the background update that adds a unique index to `T`.
|
|
|
|
This dictionary is maintained by hand.
|
|
|
|
|
|
|
|
At runtime, we constantly check to see if each of these background updates
|
|
|
|
has run. If so, we deem the coresponding table safe to upsert into, because
|
|
|
|
we can now use a native insert to do so. If not, we deem the table unsafe
|
|
|
|
to upsert into and require an emulated upsert.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-09-29 14:10:47 -04:00
|
|
|
Tables that do not appear in this dictionary are assumed to have an
|
|
|
|
appropriate unique index and therefore be safe to upsert into.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
keyvalues: The unique key columns and their new values
|
|
|
|
values: The nonunique columns and their new values
|
|
|
|
insertion_values: additional key/values to use only when inserting
|
2023-09-15 04:16:45 -04:00
|
|
|
where_clause: An index predicate to apply to the upsert.
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
Returns:
|
2021-07-22 07:39:50 -04:00
|
|
|
Returns True if a row was inserted or updated (i.e. if `values` is
|
|
|
|
not empty then this always returns True)
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values = insertion_values or {}
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
attempts = 0
|
|
|
|
while True:
|
|
|
|
try:
|
2022-09-09 06:14:10 -04:00
|
|
|
# We can autocommit if it is safe to upsert
|
|
|
|
autocommit = table not in self._unsafe_to_upsert_tables
|
2020-10-14 10:50:59 -04:00
|
|
|
|
2020-08-17 12:18:01 -04:00
|
|
|
return await self.runInteraction(
|
2019-12-04 08:52:46 -05:00
|
|
|
desc,
|
|
|
|
self.simple_upsert_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
values,
|
|
|
|
insertion_values,
|
2023-09-15 04:16:45 -04:00
|
|
|
where_clause,
|
2020-10-14 10:50:59 -04:00
|
|
|
db_autocommit=autocommit,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
2019-12-06 08:40:02 -05:00
|
|
|
except self.engine.module.IntegrityError as e:
|
2019-12-04 08:52:46 -05:00
|
|
|
attempts += 1
|
|
|
|
if attempts >= 5:
|
|
|
|
# don't retry forever, because things other than races
|
|
|
|
# can cause IntegrityErrors
|
|
|
|
raise
|
|
|
|
|
|
|
|
# presumably we raced with another transaction: let's retry.
|
|
|
|
logger.warning(
|
|
|
|
"IntegrityError when upserting into %s; retrying: %s", table, e
|
|
|
|
)
|
|
|
|
|
|
|
|
def simple_upsert_txn(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
values: Dict[str, Any],
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values: Optional[Dict[str, Any]] = None,
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause: Optional[str] = None,
|
2021-07-22 07:39:50 -04:00
|
|
|
) -> bool:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Pick the UPSERT method which works best on the platform. Either the
|
2022-09-09 06:14:10 -04:00
|
|
|
native one (Pg9.5+, SQLite >= 3.24), or fall back to an emulated method.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
|
|
|
txn: The transaction to use.
|
2020-08-20 09:00:59 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
keyvalues: The unique key tables and their new values
|
|
|
|
values: The nonunique columns and their new values
|
|
|
|
insertion_values: additional key/values to use only when inserting
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause: An index predicate to apply to the upsert.
|
2019-12-04 08:52:46 -05:00
|
|
|
Returns:
|
2021-07-22 07:39:50 -04:00
|
|
|
Returns True if a row was inserted or updated (i.e. if `values` is
|
|
|
|
not empty then this always returns True)
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values = insertion_values or {}
|
|
|
|
|
2022-09-09 06:14:10 -04:00
|
|
|
if table not in self._unsafe_to_upsert_tables:
|
2021-07-22 07:39:50 -04:00
|
|
|
return self.simple_upsert_txn_native_upsert(
|
2022-09-15 14:28:48 -04:00
|
|
|
txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
values,
|
|
|
|
insertion_values=insertion_values,
|
|
|
|
where_clause=where_clause,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
else:
|
|
|
|
return self.simple_upsert_txn_emulated(
|
|
|
|
txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
values,
|
|
|
|
insertion_values=insertion_values,
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause=where_clause,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
def simple_upsert_txn_emulated(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
values: Dict[str, Any],
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values: Optional[Dict[str, Any]] = None,
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause: Optional[str] = None,
|
2020-08-20 09:00:59 -04:00
|
|
|
lock: bool = True,
|
|
|
|
) -> bool:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
keyvalues: The unique key tables and their new values
|
|
|
|
values: The nonunique columns and their new values
|
|
|
|
insertion_values: additional key/values to use only when inserting
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause: An index predicate to apply to the upsert.
|
2020-08-20 09:00:59 -04:00
|
|
|
lock: True to lock the table when doing the upsert.
|
2022-11-28 08:42:06 -05:00
|
|
|
Must not be False unless the table has already been locked.
|
2019-12-04 08:52:46 -05:00
|
|
|
Returns:
|
2021-07-22 07:39:50 -04:00
|
|
|
Returns True if a row was inserted or updated (i.e. if `values` is
|
|
|
|
not empty then this always returns True)
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values = insertion_values or {}
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
if lock:
|
2022-11-28 08:42:06 -05:00
|
|
|
# We need to lock the table :(
|
2019-12-06 08:40:02 -05:00
|
|
|
self.engine.lock_table(txn, table)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-05-12 10:33:50 -04:00
|
|
|
def _getwhere(key: str) -> str:
|
2019-12-04 08:52:46 -05:00
|
|
|
# If the value we're passing in is None (aka NULL), we need to use
|
|
|
|
# IS, not =, as NULL = NULL equals NULL (False).
|
|
|
|
if keyvalues[key] is None:
|
|
|
|
return "%s IS ?" % (key,)
|
|
|
|
else:
|
|
|
|
return "%s = ?" % (key,)
|
|
|
|
|
2022-09-15 14:28:48 -04:00
|
|
|
# Generate a where clause of each keyvalue and optionally the provided
|
|
|
|
# index predicate.
|
|
|
|
where = [_getwhere(k) for k in keyvalues]
|
|
|
|
if where_clause:
|
|
|
|
where.append(where_clause)
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
if not values:
|
|
|
|
# If `values` is empty, then all of the values we care about are in
|
|
|
|
# the unique key, so there is nothing to UPDATE. We can just do a
|
|
|
|
# SELECT instead to see if it exists.
|
2022-09-15 14:28:48 -04:00
|
|
|
sql = "SELECT 1 FROM %s WHERE %s" % (table, " AND ".join(where))
|
2019-12-04 08:52:46 -05:00
|
|
|
sqlargs = list(keyvalues.values())
|
|
|
|
txn.execute(sql, sqlargs)
|
|
|
|
if txn.fetchall():
|
|
|
|
# We have an existing record.
|
|
|
|
return False
|
|
|
|
else:
|
|
|
|
# First try to update.
|
|
|
|
sql = "UPDATE %s SET %s WHERE %s" % (
|
|
|
|
table,
|
|
|
|
", ".join("%s = ?" % (k,) for k in values),
|
2022-09-15 14:28:48 -04:00
|
|
|
" AND ".join(where),
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
sqlargs = list(values.values()) + list(keyvalues.values())
|
|
|
|
|
|
|
|
txn.execute(sql, sqlargs)
|
|
|
|
if txn.rowcount > 0:
|
2021-07-22 07:39:50 -04:00
|
|
|
return True
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
# We didn't find any existing rows, so insert a new one
|
2021-07-15 12:46:54 -04:00
|
|
|
allvalues: Dict[str, Any] = {}
|
2019-12-04 08:52:46 -05:00
|
|
|
allvalues.update(keyvalues)
|
|
|
|
allvalues.update(values)
|
|
|
|
allvalues.update(insertion_values)
|
|
|
|
|
|
|
|
sql = "INSERT INTO %s (%s) VALUES (%s)" % (
|
|
|
|
table,
|
|
|
|
", ".join(k for k in allvalues),
|
|
|
|
", ".join("?" for _ in allvalues),
|
|
|
|
)
|
|
|
|
txn.execute(sql, list(allvalues.values()))
|
|
|
|
# successfully inserted
|
|
|
|
return True
|
|
|
|
|
|
|
|
def simple_upsert_txn_native_upsert(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
values: Dict[str, Any],
|
2021-04-08 17:38:54 -04:00
|
|
|
insertion_values: Optional[Dict[str, Any]] = None,
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause: Optional[str] = None,
|
2021-07-22 07:39:50 -04:00
|
|
|
) -> bool:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-07-22 07:39:50 -04:00
|
|
|
Use the native UPSERT functionality in PostgreSQL.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
keyvalues: The unique key tables and their new values
|
|
|
|
values: The nonunique columns and their new values
|
|
|
|
insertion_values: additional key/values to use only when inserting
|
2022-09-15 14:28:48 -04:00
|
|
|
where_clause: An index predicate to apply to the upsert.
|
2021-07-22 07:39:50 -04:00
|
|
|
|
|
|
|
Returns:
|
|
|
|
Returns True if a row was inserted or updated (i.e. if `values` is
|
|
|
|
not empty then this always returns True)
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-07-15 12:46:54 -04:00
|
|
|
allvalues: Dict[str, Any] = {}
|
2019-12-04 08:52:46 -05:00
|
|
|
allvalues.update(keyvalues)
|
2021-04-08 17:38:54 -04:00
|
|
|
allvalues.update(insertion_values or {})
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
if not values:
|
|
|
|
latter = "NOTHING"
|
|
|
|
else:
|
|
|
|
allvalues.update(values)
|
|
|
|
latter = "UPDATE SET " + ", ".join(k + "=EXCLUDED." + k for k in values)
|
|
|
|
|
2023-11-07 09:34:23 -05:00
|
|
|
sql = "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) %sDO %s" % (
|
2019-12-04 08:52:46 -05:00
|
|
|
table,
|
|
|
|
", ".join(k for k in allvalues),
|
|
|
|
", ".join("?" for _ in allvalues),
|
|
|
|
", ".join(k for k in keyvalues),
|
2023-11-07 09:34:23 -05:00
|
|
|
f"WHERE {where_clause} " if where_clause else "",
|
2019-12-04 08:52:46 -05:00
|
|
|
latter,
|
|
|
|
)
|
|
|
|
txn.execute(sql, list(allvalues.values()))
|
|
|
|
|
2021-07-22 07:39:50 -04:00
|
|
|
return bool(txn.rowcount)
|
|
|
|
|
2020-10-14 10:50:59 -04:00
|
|
|
async def simple_upsert_many(
|
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
key_names: Collection[str],
|
2021-12-13 14:01:27 -05:00
|
|
|
key_values: Collection[Collection[Any]],
|
2020-10-14 10:50:59 -04:00
|
|
|
value_names: Collection[str],
|
2021-12-13 14:01:27 -05:00
|
|
|
value_values: Collection[Collection[Any]],
|
2020-10-14 10:50:59 -04:00
|
|
|
desc: str,
|
|
|
|
) -> None:
|
|
|
|
"""
|
|
|
|
Upsert, many times.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
table: The table to upsert into
|
|
|
|
key_names: The key column names.
|
|
|
|
key_values: A list of each row's key column values.
|
|
|
|
value_names: The value column names
|
|
|
|
value_values: A list of each row's value column values.
|
|
|
|
Ignored if value_names is empty.
|
|
|
|
"""
|
|
|
|
|
2022-09-09 06:14:10 -04:00
|
|
|
# We can autocommit if it safe to upsert
|
|
|
|
autocommit = table not in self._unsafe_to_upsert_tables
|
2020-10-14 10:50:59 -04:00
|
|
|
|
2022-04-08 10:29:13 -04:00
|
|
|
await self.runInteraction(
|
2020-10-14 10:50:59 -04:00
|
|
|
desc,
|
|
|
|
self.simple_upsert_many_txn,
|
|
|
|
table,
|
|
|
|
key_names,
|
|
|
|
key_values,
|
|
|
|
value_names,
|
|
|
|
value_values,
|
|
|
|
db_autocommit=autocommit,
|
|
|
|
)
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
def simple_upsert_many_txn(
|
2020-05-05 20:08:15 -04:00
|
|
|
self,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
key_names: Collection[str],
|
|
|
|
key_values: Collection[Iterable[Any]],
|
|
|
|
value_names: Collection[str],
|
2023-11-07 14:00:25 -05:00
|
|
|
value_values: Collection[Iterable[Any]],
|
2020-05-05 20:08:15 -04:00
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Upsert, many times.
|
|
|
|
|
|
|
|
Args:
|
2020-05-05 20:08:15 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
key_names: The key column names.
|
|
|
|
key_values: A list of each row's key column values.
|
|
|
|
value_names: The value column names
|
|
|
|
value_values: A list of each row's value column values.
|
|
|
|
Ignored if value_names is empty.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2023-11-07 14:00:25 -05:00
|
|
|
# If there's nothing to upsert, then skip executing the query.
|
|
|
|
if not key_values:
|
|
|
|
return
|
|
|
|
|
|
|
|
# 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))]
|
|
|
|
elif 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."
|
|
|
|
)
|
|
|
|
|
2022-09-09 06:14:10 -04:00
|
|
|
if table not in self._unsafe_to_upsert_tables:
|
2019-12-04 08:52:46 -05:00
|
|
|
return self.simple_upsert_many_txn_native_upsert(
|
|
|
|
txn, table, key_names, key_values, value_names, value_values
|
|
|
|
)
|
|
|
|
else:
|
|
|
|
return self.simple_upsert_many_txn_emulated(
|
2022-11-28 08:42:06 -05:00
|
|
|
txn,
|
|
|
|
table,
|
|
|
|
key_names,
|
|
|
|
key_values,
|
|
|
|
value_names,
|
|
|
|
value_values,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
def simple_upsert_many_txn_emulated(
|
2020-05-05 20:08:15 -04:00
|
|
|
self,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
key_names: Iterable[str],
|
|
|
|
key_values: Collection[Iterable[Any]],
|
|
|
|
value_names: Collection[str],
|
2020-09-04 07:22:23 -04:00
|
|
|
value_values: Iterable[Iterable[Any]],
|
2020-05-05 20:08:15 -04:00
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Upsert, many times, but without native UPSERT support or batching.
|
|
|
|
|
|
|
|
Args:
|
2020-05-05 20:08:15 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
key_names: The key column names.
|
|
|
|
key_values: A list of each row's key column values.
|
|
|
|
value_names: The value column names
|
|
|
|
value_values: A list of each row's value column values.
|
|
|
|
Ignored if value_names is empty.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
|
2022-11-28 08:42:06 -05:00
|
|
|
# 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.
|
2023-06-16 10:25:44 -04:00
|
|
|
self.engine.lock_table(txn, table)
|
2022-04-08 10:29:13 -04:00
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
for keyv, valv in zip(key_values, value_values):
|
2023-03-28 04:46:47 -04:00
|
|
|
_keys = dict(zip(key_names, keyv))
|
|
|
|
_vals = dict(zip(value_names, valv))
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-04-08 10:29:13 -04:00
|
|
|
self.simple_upsert_txn_emulated(txn, table, _keys, _vals, lock=False)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
def simple_upsert_many_txn_native_upsert(
|
2020-05-05 20:08:15 -04:00
|
|
|
self,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
key_names: Collection[str],
|
|
|
|
key_values: Collection[Iterable[Any]],
|
|
|
|
value_names: Collection[str],
|
|
|
|
value_values: Iterable[Iterable[Any]],
|
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Upsert, many times, using batching where possible.
|
|
|
|
|
|
|
|
Args:
|
2020-05-05 20:08:15 -04:00
|
|
|
table: The table to upsert into
|
|
|
|
key_names: The key column names.
|
|
|
|
key_values: A list of each row's key column values.
|
|
|
|
value_names: The value column names
|
|
|
|
value_values: A list of each row's value column values.
|
|
|
|
Ignored if value_names is empty.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-07-15 12:46:54 -04:00
|
|
|
allnames: List[str] = []
|
2019-12-04 08:52:46 -05:00
|
|
|
allnames.extend(key_names)
|
|
|
|
allnames.extend(value_names)
|
|
|
|
|
|
|
|
if not value_names:
|
|
|
|
latter = "NOTHING"
|
|
|
|
else:
|
|
|
|
latter = "UPDATE SET " + ", ".join(
|
|
|
|
k + "=EXCLUDED." + k for k in value_names
|
|
|
|
)
|
|
|
|
|
|
|
|
args = []
|
|
|
|
|
|
|
|
for x, y in zip(key_values, value_values):
|
|
|
|
args.append(tuple(x) + tuple(y))
|
|
|
|
|
2021-09-03 11:35:49 -04:00
|
|
|
if isinstance(txn.database_engine, PostgresEngine):
|
|
|
|
# We use `execute_values` as it can be a lot faster than `execute_batch`,
|
|
|
|
# but it's only available on postgres.
|
|
|
|
sql = "INSERT INTO %s (%s) VALUES ? ON CONFLICT (%s) DO %s" % (
|
|
|
|
table,
|
|
|
|
", ".join(k for k in allnames),
|
|
|
|
", ".join(key_names),
|
|
|
|
latter,
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute_values(sql, args, fetch=False)
|
|
|
|
|
|
|
|
else:
|
|
|
|
sql = "INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO %s" % (
|
|
|
|
table,
|
|
|
|
", ".join(k for k in allnames),
|
|
|
|
", ".join("?" for _ in allnames),
|
|
|
|
", ".join(key_names),
|
|
|
|
latter,
|
|
|
|
)
|
|
|
|
|
|
|
|
return txn.execute_batch(sql, args)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-26 07:19:32 -04:00
|
|
|
@overload
|
|
|
|
async def simple_select_one(
|
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2021-12-15 12:00:50 -05:00
|
|
|
retcols: Collection[str],
|
2020-08-26 07:19:32 -04:00
|
|
|
allow_none: Literal[False] = False,
|
|
|
|
desc: str = "simple_select_one",
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Tuple[Any, ...]: ...
|
2020-08-26 07:19:32 -04:00
|
|
|
|
|
|
|
@overload
|
|
|
|
async def simple_select_one(
|
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2021-12-15 12:00:50 -05:00
|
|
|
retcols: Collection[str],
|
2020-08-26 07:19:32 -04:00
|
|
|
allow_none: Literal[True] = True,
|
|
|
|
desc: str = "simple_select_one",
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Optional[Tuple[Any, ...]]: ...
|
2020-08-26 07:19:32 -04:00
|
|
|
|
|
|
|
async def simple_select_one(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2021-12-15 12:00:50 -05:00
|
|
|
retcols: Collection[str],
|
2020-08-20 09:00:59 -04:00
|
|
|
allow_none: bool = False,
|
|
|
|
desc: str = "simple_select_one",
|
2023-11-09 11:13:31 -05:00
|
|
|
) -> Optional[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which is expected to
|
|
|
|
return a single row, returning multiple columns from it.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
|
|
|
retcols: list of strings giving the names of the columns to return
|
|
|
|
allow_none: If true, return None instead of failing if the SELECT
|
|
|
|
statement returns no rows
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-08-26 07:19:32 -04:00
|
|
|
return await self.runInteraction(
|
2020-10-14 10:50:59 -04:00
|
|
|
desc,
|
|
|
|
self.simple_select_one_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
retcols,
|
|
|
|
allow_none,
|
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
2020-09-02 12:52:38 -04:00
|
|
|
@overload
|
|
|
|
async def simple_select_one_onecol(
|
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2020-09-04 07:02:29 -04:00
|
|
|
retcol: str,
|
2020-09-02 12:52:38 -04:00
|
|
|
allow_none: Literal[False] = False,
|
|
|
|
desc: str = "simple_select_one_onecol",
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Any: ...
|
2020-09-02 12:52:38 -04:00
|
|
|
|
|
|
|
@overload
|
|
|
|
async def simple_select_one_onecol(
|
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2020-09-04 07:02:29 -04:00
|
|
|
retcol: str,
|
2020-09-02 12:52:38 -04:00
|
|
|
allow_none: Literal[True] = True,
|
|
|
|
desc: str = "simple_select_one_onecol",
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Optional[Any]: ...
|
2020-09-02 12:52:38 -04:00
|
|
|
|
2020-08-26 07:19:32 -04:00
|
|
|
async def simple_select_one_onecol(
|
2019-12-04 08:52:46 -05:00
|
|
|
self,
|
2020-08-20 09:00:59 -04:00
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2020-09-04 07:02:29 -04:00
|
|
|
retcol: str,
|
2020-08-20 09:00:59 -04:00
|
|
|
allow_none: bool = False,
|
|
|
|
desc: str = "simple_select_one_onecol",
|
2020-08-26 07:19:32 -04:00
|
|
|
) -> Optional[Any]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which is expected to
|
|
|
|
return a single row, returning a single column from it.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
|
|
|
retcol: string giving the name of the column to return
|
2022-10-18 07:33:18 -04:00
|
|
|
allow_none: If true, return None instead of raising StoreError if the SELECT
|
2020-08-20 09:00:59 -04:00
|
|
|
statement returns no rows
|
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-08-26 07:19:32 -04:00
|
|
|
return await self.runInteraction(
|
2019-12-04 08:52:46 -05:00
|
|
|
desc,
|
|
|
|
self.simple_select_one_onecol_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
retcol,
|
|
|
|
allow_none=allow_none,
|
2020-10-14 10:50:59 -04:00
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
2020-09-02 15:03:12 -04:00
|
|
|
@overload
|
|
|
|
@classmethod
|
|
|
|
def simple_select_one_onecol_txn(
|
|
|
|
cls,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2020-09-04 07:02:29 -04:00
|
|
|
retcol: str,
|
2020-09-02 15:03:12 -04:00
|
|
|
allow_none: Literal[False] = False,
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Any: ...
|
2020-09-02 15:03:12 -04:00
|
|
|
|
|
|
|
@overload
|
|
|
|
@classmethod
|
|
|
|
def simple_select_one_onecol_txn(
|
|
|
|
cls,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2020-09-04 07:02:29 -04:00
|
|
|
retcol: str,
|
2020-09-02 15:03:12 -04:00
|
|
|
allow_none: Literal[True] = True,
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Optional[Any]: ...
|
2020-09-02 15:03:12 -04:00
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
@classmethod
|
|
|
|
def simple_select_one_onecol_txn(
|
2020-08-20 09:00:59 -04:00
|
|
|
cls,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2020-09-04 07:02:29 -04:00
|
|
|
retcol: str,
|
2020-08-20 09:00:59 -04:00
|
|
|
allow_none: bool = False,
|
|
|
|
) -> Optional[Any]:
|
2019-12-04 08:52:46 -05:00
|
|
|
ret = cls.simple_select_onecol_txn(
|
|
|
|
txn, table=table, keyvalues=keyvalues, retcol=retcol
|
|
|
|
)
|
|
|
|
|
|
|
|
if ret:
|
|
|
|
return ret[0]
|
|
|
|
else:
|
|
|
|
if allow_none:
|
|
|
|
return None
|
|
|
|
else:
|
|
|
|
raise StoreError(404, "No row found")
|
|
|
|
|
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_select_onecol_txn(
|
2021-02-16 17:32:34 -05:00
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
retcol: str,
|
2020-08-20 09:00:59 -04:00
|
|
|
) -> List[Any]:
|
2019-12-04 08:52:46 -05:00
|
|
|
sql = ("SELECT %(retcol)s FROM %(table)s") % {"retcol": retcol, "table": table}
|
|
|
|
|
|
|
|
if keyvalues:
|
2020-06-15 07:03:36 -04:00
|
|
|
sql += " WHERE %s" % " AND ".join("%s = ?" % k for k in keyvalues.keys())
|
2019-12-04 08:52:46 -05:00
|
|
|
txn.execute(sql, list(keyvalues.values()))
|
|
|
|
else:
|
|
|
|
txn.execute(sql)
|
|
|
|
|
|
|
|
return [r[0] for r in txn]
|
|
|
|
|
2020-08-27 07:08:38 -04:00
|
|
|
async def simple_select_onecol(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Optional[Dict[str, Any]],
|
|
|
|
retcol: str,
|
|
|
|
desc: str = "simple_select_onecol",
|
2020-08-27 07:08:38 -04:00
|
|
|
) -> List[Any]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which returns a list
|
|
|
|
comprising of the values of the named column from the selected rows.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: table name
|
|
|
|
keyvalues: column names and values to select the rows with
|
|
|
|
retcol: column whos value we wish to retrieve.
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Returns:
|
2020-08-27 07:08:38 -04:00
|
|
|
Results in a list
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-08-27 07:08:38 -04:00
|
|
|
return await self.runInteraction(
|
2020-10-14 10:50:59 -04:00
|
|
|
desc,
|
|
|
|
self.simple_select_onecol_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
retcol,
|
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
2020-08-27 07:08:38 -04:00
|
|
|
async def simple_select_list(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Optional[Dict[str, Any]],
|
2021-12-15 12:00:50 -05:00
|
|
|
retcols: Collection[str],
|
2020-08-20 09:00:59 -04:00
|
|
|
desc: str = "simple_select_list",
|
2023-10-26 13:01:36 -04:00
|
|
|
) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which may return zero or
|
2023-10-26 13:01:36 -04:00
|
|
|
more rows, returning the result as a list of tuples.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: the table name
|
|
|
|
keyvalues:
|
2019-12-04 08:52:46 -05:00
|
|
|
column names and values to select the rows with, or None to not
|
|
|
|
apply a WHERE clause.
|
2020-08-20 09:00:59 -04:00
|
|
|
retcols: the names of the columns to return
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2020-08-27 07:08:38 -04:00
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
Returns:
|
2023-10-26 13:01:36 -04:00
|
|
|
A list of tuples, one per result row, each the retcolumn's value for the row.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-08-27 07:08:38 -04:00
|
|
|
return await self.runInteraction(
|
2020-10-14 10:50:59 -04:00
|
|
|
desc,
|
|
|
|
self.simple_select_list_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
retcols,
|
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
@classmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_select_list_txn(
|
|
|
|
cls,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Optional[Dict[str, Any]],
|
|
|
|
retcols: Iterable[str],
|
2023-10-26 13:01:36 -04:00
|
|
|
) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which may return zero or
|
2023-10-26 13:01:36 -04:00
|
|
|
more rows, returning the result as a list of tuples.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
txn: Transaction object
|
|
|
|
table: the table name
|
|
|
|
keyvalues:
|
2019-12-04 08:52:46 -05:00
|
|
|
column names and values to select the rows with, or None to not
|
|
|
|
apply a WHERE clause.
|
2020-08-20 09:00:59 -04:00
|
|
|
retcols: the names of the columns to return
|
2022-12-31 22:40:46 -05:00
|
|
|
|
|
|
|
Returns:
|
2023-10-26 13:01:36 -04:00
|
|
|
A list of tuples, one per result row, each the retcolumn's value for the row.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
if keyvalues:
|
|
|
|
sql = "SELECT %s FROM %s WHERE %s" % (
|
|
|
|
", ".join(retcols),
|
|
|
|
table,
|
|
|
|
" AND ".join("%s = ?" % (k,) for k in keyvalues),
|
|
|
|
)
|
|
|
|
txn.execute(sql, list(keyvalues.values()))
|
|
|
|
else:
|
|
|
|
sql = "SELECT %s FROM %s" % (", ".join(retcols), table)
|
|
|
|
txn.execute(sql)
|
|
|
|
|
2023-10-26 13:01:36 -04:00
|
|
|
return txn.fetchall()
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-17 12:18:01 -04:00
|
|
|
async def simple_select_many_batch(
|
2019-12-04 08:52:46 -05:00
|
|
|
self,
|
2020-08-20 09:00:59 -04:00
|
|
|
table: str,
|
|
|
|
column: str,
|
|
|
|
iterable: Iterable[Any],
|
2021-12-15 12:00:50 -05:00
|
|
|
retcols: Collection[str],
|
2021-04-08 17:38:54 -04:00
|
|
|
keyvalues: Optional[Dict[str, Any]] = None,
|
2020-08-20 09:00:59 -04:00
|
|
|
desc: str = "simple_select_many_batch",
|
|
|
|
batch_size: int = 100,
|
2023-10-11 13:24:56 -04:00
|
|
|
) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which may return zero or
|
2023-10-11 13:24:56 -04:00
|
|
|
more rows.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-27 14:16:41 -04:00
|
|
|
Filters rows by whether the value of `column` is in `iterable`.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
column: column name to test for inclusion against `iterable`
|
|
|
|
iterable: list
|
|
|
|
retcols: list of strings giving the names of the columns to return
|
2020-08-27 14:16:41 -04:00
|
|
|
keyvalues: dict of column names and values to select the rows with
|
|
|
|
desc: description of the transaction, for logging and metrics
|
|
|
|
batch_size: the number of rows for each select query
|
2023-10-11 13:24:56 -04:00
|
|
|
|
|
|
|
Returns:
|
|
|
|
The results as a list of tuples.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2021-04-08 17:38:54 -04:00
|
|
|
keyvalues = keyvalues or {}
|
|
|
|
|
2023-10-11 13:24:56 -04:00
|
|
|
results: List[Tuple[Any, ...]] = []
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2021-12-15 12:00:50 -05:00
|
|
|
for chunk in batch_iter(iterable, batch_size):
|
2020-08-17 12:18:01 -04:00
|
|
|
rows = await self.runInteraction(
|
2019-12-04 08:52:46 -05:00
|
|
|
desc,
|
|
|
|
self.simple_select_many_txn,
|
|
|
|
table,
|
|
|
|
column,
|
|
|
|
chunk,
|
|
|
|
keyvalues,
|
|
|
|
retcols,
|
2020-10-14 10:50:59 -04:00
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
results.extend(rows)
|
|
|
|
|
|
|
|
return results
|
|
|
|
|
|
|
|
@classmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_select_many_txn(
|
|
|
|
cls,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
column: str,
|
2021-09-20 05:26:13 -04:00
|
|
|
iterable: Collection[Any],
|
2020-08-20 09:00:59 -04:00
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
retcols: Iterable[str],
|
2023-10-11 13:24:56 -04:00
|
|
|
) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a SELECT query on the named table, which may return zero or
|
2023-10-11 13:24:56 -04:00
|
|
|
more rows.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-27 14:16:41 -04:00
|
|
|
Filters rows by whether the value of `column` is in `iterable`.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
txn: Transaction object
|
|
|
|
table: string giving the table name
|
|
|
|
column: column name to test for inclusion against `iterable`
|
|
|
|
iterable: list
|
|
|
|
keyvalues: dict of column names and values to select the rows with
|
|
|
|
retcols: list of strings giving the names of the columns to return
|
2023-10-11 13:24:56 -04:00
|
|
|
|
|
|
|
Returns:
|
|
|
|
The results as a list of tuples.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2023-11-07 14:00:25 -05:00
|
|
|
# If there's nothing to select, then skip executing the query.
|
2019-12-04 08:52:46 -05:00
|
|
|
if not iterable:
|
|
|
|
return []
|
|
|
|
|
|
|
|
clause, values = make_in_list_sql_clause(txn.database_engine, column, iterable)
|
|
|
|
clauses = [clause]
|
|
|
|
|
2020-06-15 07:03:36 -04:00
|
|
|
for key, value in keyvalues.items():
|
2019-12-04 08:52:46 -05:00
|
|
|
clauses.append("%s = ?" % (key,))
|
|
|
|
values.append(value)
|
|
|
|
|
|
|
|
sql = "SELECT %s FROM %s WHERE %s" % (
|
|
|
|
", ".join(retcols),
|
|
|
|
table,
|
|
|
|
" AND ".join(clauses),
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute(sql, values)
|
2023-10-11 13:24:56 -04:00
|
|
|
return txn.fetchall()
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-27 07:08:38 -04:00
|
|
|
async def simple_update(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
updatevalues: Dict[str, Any],
|
|
|
|
desc: str,
|
2020-08-27 07:08:38 -04:00
|
|
|
) -> int:
|
2022-12-31 22:40:46 -05:00
|
|
|
"""
|
|
|
|
Update rows in the given database table.
|
|
|
|
If the given keyvalues don't match anything, nothing will be updated.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
table: The database table to update.
|
|
|
|
keyvalues: A mapping of column name to value to match rows on.
|
|
|
|
updatevalues: A mapping of column name to value to replace in any matched rows.
|
|
|
|
desc: description of the transaction, for logging and metrics.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
The number of rows that were updated. Will be 0 if no matching rows were found.
|
|
|
|
"""
|
2020-08-27 07:08:38 -04:00
|
|
|
return await self.runInteraction(
|
2019-12-04 08:52:46 -05:00
|
|
|
desc, self.simple_update_txn, table, keyvalues, updatevalues
|
|
|
|
)
|
|
|
|
|
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_update_txn(
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
updatevalues: Dict[str, Any],
|
|
|
|
) -> int:
|
2022-12-31 22:40:46 -05:00
|
|
|
"""
|
|
|
|
Update rows in the given database table.
|
|
|
|
If the given keyvalues don't match anything, nothing will be updated.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
txn: The database transaction object.
|
|
|
|
table: The database table to update.
|
|
|
|
keyvalues: A mapping of column name to value to match rows on.
|
|
|
|
updatevalues: A mapping of column name to value to replace in any matched rows.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
The number of rows that were updated. Will be 0 if no matching rows were found.
|
|
|
|
"""
|
2019-12-04 08:52:46 -05:00
|
|
|
if keyvalues:
|
2020-06-15 07:03:36 -04:00
|
|
|
where = "WHERE %s" % " AND ".join("%s = ?" % k for k in keyvalues.keys())
|
2019-12-04 08:52:46 -05:00
|
|
|
else:
|
|
|
|
where = ""
|
|
|
|
|
|
|
|
update_sql = "UPDATE %s SET %s %s" % (
|
|
|
|
table,
|
|
|
|
", ".join("%s = ?" % (k,) for k in updatevalues),
|
|
|
|
where,
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute(update_sql, list(updatevalues.values()) + list(keyvalues.values()))
|
|
|
|
|
|
|
|
return txn.rowcount
|
|
|
|
|
2022-04-08 10:29:13 -04:00
|
|
|
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."
|
|
|
|
)
|
2023-11-07 14:00:25 -05:00
|
|
|
# If there is nothing to update, then skip executing the query.
|
|
|
|
if not key_values:
|
|
|
|
return
|
2022-04-08 10:29:13 -04:00
|
|
|
|
|
|
|
# List of tuples of (value values, then key values)
|
|
|
|
# (This matches the order needed for the query)
|
2023-11-07 14:02:09 -05:00
|
|
|
args = [tuple(vv) + tuple(kv) for vv, kv in zip(value_values, key_values)]
|
2022-04-08 10:29:13 -04:00
|
|
|
|
|
|
|
# '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 = ?
|
2023-11-07 09:34:23 -05:00
|
|
|
sql = f"UPDATE {table} SET {set_clause} {where_clause}"
|
2022-04-08 10:29:13 -04:00
|
|
|
|
|
|
|
txn.execute_batch(sql, args)
|
|
|
|
|
2020-08-27 07:08:38 -04:00
|
|
|
async def simple_update_one(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
updatevalues: Dict[str, Any],
|
|
|
|
desc: str = "simple_update_one",
|
2020-08-27 07:08:38 -04:00
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes an UPDATE query on the named table, setting new values for
|
|
|
|
columns in a row matching the key values.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
|
|
|
updatevalues: dict giving column names and values to update
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-08-27 07:08:38 -04:00
|
|
|
await self.runInteraction(
|
2020-10-14 10:50:59 -04:00
|
|
|
desc,
|
|
|
|
self.simple_update_one_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
updatevalues,
|
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
@classmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_update_one_txn(
|
|
|
|
cls,
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
updatevalues: Dict[str, Any],
|
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
rowcount = cls.simple_update_txn(txn, table, keyvalues, updatevalues)
|
|
|
|
|
|
|
|
if rowcount == 0:
|
|
|
|
raise StoreError(404, "No row found (%s)" % (table,))
|
|
|
|
if rowcount > 1:
|
|
|
|
raise StoreError(500, "More than one row matched (%s)" % (table,))
|
|
|
|
|
2020-08-20 09:00:59 -04:00
|
|
|
# Ideally we could use the overload decorator here to specify that the
|
|
|
|
# return type is only optional if allow_none is True, but this does not work
|
|
|
|
# when you call a static method from an instance.
|
|
|
|
# See https://github.com/python/mypy/issues/7781
|
2019-12-04 08:52:46 -05:00
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_select_one_txn(
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keyvalues: Dict[str, Any],
|
2021-12-15 12:00:50 -05:00
|
|
|
retcols: Collection[str],
|
2020-08-20 09:00:59 -04:00
|
|
|
allow_none: bool = False,
|
2023-11-09 11:13:31 -05:00
|
|
|
) -> Optional[Tuple[Any, ...]]:
|
2022-11-22 11:46:52 -05:00
|
|
|
select_sql = "SELECT %s FROM %s" % (", ".join(retcols), table)
|
|
|
|
|
|
|
|
if keyvalues:
|
|
|
|
select_sql += " WHERE %s" % (" AND ".join("%s = ?" % k for k in keyvalues),)
|
|
|
|
txn.execute(select_sql, list(keyvalues.values()))
|
|
|
|
else:
|
|
|
|
txn.execute(select_sql)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
row = txn.fetchone()
|
|
|
|
|
|
|
|
if not row:
|
|
|
|
if allow_none:
|
|
|
|
return None
|
|
|
|
raise StoreError(404, "No row found (%s)" % (table,))
|
|
|
|
if txn.rowcount > 1:
|
|
|
|
raise StoreError(500, "More than one row matched (%s)" % (table,))
|
|
|
|
|
2023-11-09 11:13:31 -05:00
|
|
|
return row
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-08-27 07:41:01 -04:00
|
|
|
async def simple_delete_one(
|
2020-08-20 09:00:59 -04:00
|
|
|
self, table: str, keyvalues: Dict[str, Any], desc: str = "simple_delete_one"
|
2020-08-27 07:41:01 -04:00
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a DELETE query on the named table, expecting to delete a
|
|
|
|
single row.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
2020-08-27 14:16:41 -04:00
|
|
|
desc: description of the transaction, for logging and metrics
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2020-10-14 10:50:59 -04:00
|
|
|
await self.runInteraction(
|
2021-02-16 17:32:34 -05:00
|
|
|
desc,
|
|
|
|
self.simple_delete_one_txn,
|
|
|
|
table,
|
|
|
|
keyvalues,
|
|
|
|
db_autocommit=True,
|
2020-10-14 10:50:59 -04:00
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_delete_one_txn(
|
|
|
|
txn: LoggingTransaction, table: str, keyvalues: Dict[str, Any]
|
|
|
|
) -> None:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a DELETE query on the named table, expecting to delete a
|
|
|
|
single row.
|
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
sql = "DELETE FROM %s WHERE %s" % (
|
|
|
|
table,
|
|
|
|
" AND ".join("%s = ?" % (k,) for k in keyvalues),
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute(sql, list(keyvalues.values()))
|
|
|
|
if txn.rowcount == 0:
|
|
|
|
raise StoreError(404, "No row found (%s)" % (table,))
|
|
|
|
if txn.rowcount > 1:
|
|
|
|
raise StoreError(500, "More than one row matched (%s)" % (table,))
|
|
|
|
|
2020-08-27 14:16:41 -04:00
|
|
|
async def simple_delete(
|
|
|
|
self, table: str, keyvalues: Dict[str, Any], desc: str
|
|
|
|
) -> int:
|
|
|
|
"""Executes a DELETE query on the named table.
|
|
|
|
|
|
|
|
Filters rows by the key-value pairs.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
|
|
|
desc: description of the transaction, for logging and metrics
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
The number of deleted rows.
|
|
|
|
"""
|
2020-10-14 10:50:59 -04:00
|
|
|
return await self.runInteraction(
|
|
|
|
desc, self.simple_delete_txn, table, keyvalues, db_autocommit=True
|
|
|
|
)
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_delete_txn(
|
|
|
|
txn: LoggingTransaction, table: str, keyvalues: Dict[str, Any]
|
|
|
|
) -> int:
|
2020-08-27 14:16:41 -04:00
|
|
|
"""Executes a DELETE query on the named table.
|
|
|
|
|
|
|
|
Filters rows by the key-value pairs.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
table: string giving the table name
|
|
|
|
keyvalues: dict of column names and values to select the row with
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
The number of deleted rows.
|
|
|
|
"""
|
2019-12-04 08:52:46 -05:00
|
|
|
sql = "DELETE FROM %s WHERE %s" % (
|
|
|
|
table,
|
|
|
|
" AND ".join("%s = ?" % (k,) for k in keyvalues),
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute(sql, list(keyvalues.values()))
|
|
|
|
return txn.rowcount
|
|
|
|
|
2020-08-27 07:41:01 -04:00
|
|
|
async def simple_delete_many(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
|
|
|
table: str,
|
|
|
|
column: str,
|
2021-12-13 14:01:27 -05:00
|
|
|
iterable: Collection[Any],
|
2020-08-20 09:00:59 -04:00
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
desc: str,
|
2020-08-27 07:41:01 -04:00
|
|
|
) -> int:
|
2020-08-27 14:16:41 -04:00
|
|
|
"""Executes a DELETE query on the named table.
|
|
|
|
|
|
|
|
Filters rows by if value of `column` is in `iterable`.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
table: string giving the table name
|
|
|
|
column: column name to test for inclusion against `iterable`
|
2021-12-13 14:01:27 -05:00
|
|
|
iterable: list of values to match against `column`. NB cannot be a generator
|
|
|
|
as it may be evaluated multiple times.
|
2020-08-27 14:16:41 -04:00
|
|
|
keyvalues: dict of column names and values to select the rows with
|
|
|
|
desc: description of the transaction, for logging and metrics
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Number rows deleted
|
|
|
|
"""
|
2020-08-27 07:41:01 -04:00
|
|
|
return await self.runInteraction(
|
2020-10-14 10:50:59 -04:00
|
|
|
desc,
|
|
|
|
self.simple_delete_many_txn,
|
|
|
|
table,
|
|
|
|
column,
|
|
|
|
iterable,
|
|
|
|
keyvalues,
|
|
|
|
db_autocommit=True,
|
2019-12-04 08:52:46 -05:00
|
|
|
)
|
|
|
|
|
|
|
|
@staticmethod
|
2020-08-20 09:00:59 -04:00
|
|
|
def simple_delete_many_txn(
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
column: str,
|
2021-09-20 05:26:13 -04:00
|
|
|
values: Collection[Any],
|
2020-08-20 09:00:59 -04:00
|
|
|
keyvalues: Dict[str, Any],
|
|
|
|
) -> int:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Executes a DELETE query on the named table.
|
|
|
|
|
2021-09-20 05:26:13 -04:00
|
|
|
Deletes the rows:
|
|
|
|
- whose value of `column` is in `values`; AND
|
|
|
|
- that match extra column-value pairs specified in `keyvalues`.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
txn: Transaction object
|
|
|
|
table: string giving the table name
|
2021-09-20 05:26:13 -04:00
|
|
|
column: column name to test for inclusion against `values`
|
|
|
|
values: values of `column` which choose rows to delete
|
|
|
|
keyvalues: dict of extra column names and values to select the rows
|
|
|
|
with. They will be ANDed together with the main predicate.
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Returns:
|
2020-08-20 09:00:59 -04:00
|
|
|
Number rows deleted
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
2023-11-07 14:00:25 -05:00
|
|
|
# If there's nothing to delete, then skip executing the query.
|
2021-09-20 05:26:13 -04:00
|
|
|
if not values:
|
2019-12-04 08:52:46 -05:00
|
|
|
return 0
|
|
|
|
|
2021-09-20 05:26:13 -04:00
|
|
|
clause, values = make_in_list_sql_clause(txn.database_engine, column, values)
|
2019-12-04 08:52:46 -05:00
|
|
|
clauses = [clause]
|
|
|
|
|
2020-06-15 07:03:36 -04:00
|
|
|
for key, value in keyvalues.items():
|
2019-12-04 08:52:46 -05:00
|
|
|
clauses.append("%s = ?" % (key,))
|
|
|
|
values.append(value)
|
|
|
|
|
2023-11-07 09:34:23 -05:00
|
|
|
sql = "DELETE FROM %s WHERE %s" % (table, " AND ".join(clauses))
|
2019-12-04 08:52:46 -05:00
|
|
|
txn.execute(sql, values)
|
|
|
|
|
|
|
|
return txn.rowcount
|
|
|
|
|
2023-07-05 05:43:19 -04:00
|
|
|
@staticmethod
|
|
|
|
def simple_delete_many_batch_txn(
|
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
keys: Collection[str],
|
|
|
|
values: Iterable[Iterable[Any]],
|
|
|
|
) -> None:
|
|
|
|
"""Executes a DELETE query on the named table.
|
|
|
|
|
|
|
|
The input is given as a list of rows, where each row is a list of values.
|
|
|
|
(Actually any iterable is fine.)
|
|
|
|
|
|
|
|
Args:
|
|
|
|
txn: The transaction to use.
|
|
|
|
table: string giving the table name
|
|
|
|
keys: list of column names
|
|
|
|
values: for each row, a list of values in the same order as `keys`
|
|
|
|
"""
|
|
|
|
|
|
|
|
if isinstance(txn.database_engine, PostgresEngine):
|
|
|
|
# We use `execute_values` as it can be a lot faster than `execute_batch`,
|
|
|
|
# but it's only available on postgres.
|
|
|
|
sql = "DELETE FROM %s WHERE (%s) IN (VALUES ?)" % (
|
|
|
|
table,
|
|
|
|
", ".join(k for k in keys),
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute_values(sql, values, fetch=False)
|
|
|
|
else:
|
|
|
|
sql = "DELETE FROM %s WHERE (%s) = (%s)" % (
|
|
|
|
table,
|
|
|
|
", ".join(k for k in keys),
|
|
|
|
", ".join("?" for _ in keys),
|
|
|
|
)
|
|
|
|
|
|
|
|
txn.execute_batch(sql, values)
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
def get_cache_dict(
|
2020-08-20 09:00:59 -04:00
|
|
|
self,
|
2020-10-02 10:20:45 -04:00
|
|
|
db_conn: LoggingDatabaseConnection,
|
2020-08-20 09:00:59 -04:00
|
|
|
table: str,
|
|
|
|
entity_column: str,
|
|
|
|
stream_column: str,
|
|
|
|
max_value: int,
|
|
|
|
limit: int = 100000,
|
|
|
|
) -> Tuple[Dict[Any, int], int]:
|
2022-04-05 09:26:41 -04:00
|
|
|
"""Gets roughly the last N changes in the given stream table as a
|
|
|
|
map from entity to the stream ID of the most recent change.
|
|
|
|
|
|
|
|
Also returns the minimum stream ID.
|
|
|
|
"""
|
|
|
|
|
|
|
|
# This may return many rows for the same entity, but the `limit` is only
|
|
|
|
# a suggestion so we don't care that much.
|
|
|
|
#
|
|
|
|
# Note: Some stream tables can have multiple rows with the same stream
|
|
|
|
# ID. Instead of handling this with complicated SQL, we instead simply
|
|
|
|
# add one to the returned minimum stream ID to ensure correctness.
|
|
|
|
sql = f"""
|
|
|
|
SELECT {entity_column}, {stream_column}
|
|
|
|
FROM {table}
|
|
|
|
ORDER BY {stream_column} DESC
|
|
|
|
LIMIT ?
|
|
|
|
"""
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2020-10-02 10:20:45 -04:00
|
|
|
txn = db_conn.cursor(txn_name="get_cache_dict")
|
2022-04-05 09:26:41 -04:00
|
|
|
txn.execute(sql, (limit,))
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2022-04-05 09:26:41 -04:00
|
|
|
# The rows come out in reverse stream ID order, so we want to keep the
|
|
|
|
# stream ID of the first row for each entity.
|
|
|
|
cache: Dict[Any, int] = {}
|
|
|
|
for row in txn:
|
|
|
|
cache.setdefault(row[0], int(row[1]))
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
txn.close()
|
|
|
|
|
|
|
|
if cache:
|
2022-04-05 09:26:41 -04:00
|
|
|
# We add one here as we don't know if we have all rows for the
|
|
|
|
# minimum stream ID.
|
|
|
|
min_val = min(cache.values()) + 1
|
2019-12-04 08:52:46 -05:00
|
|
|
else:
|
|
|
|
min_val = max_value
|
|
|
|
|
|
|
|
return cache, min_val
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
def simple_select_list_paginate_txn(
|
|
|
|
cls,
|
2020-08-20 09:00:59 -04:00
|
|
|
txn: LoggingTransaction,
|
|
|
|
table: str,
|
|
|
|
orderby: str,
|
|
|
|
start: int,
|
|
|
|
limit: int,
|
|
|
|
retcols: Iterable[str],
|
|
|
|
filters: Optional[Dict[str, Any]] = None,
|
|
|
|
keyvalues: Optional[Dict[str, Any]] = None,
|
2021-03-25 06:34:23 -04:00
|
|
|
exclude_keyvalues: Optional[Dict[str, Any]] = None,
|
2020-08-20 09:00:59 -04:00
|
|
|
order_direction: str = "ASC",
|
2023-10-06 11:41:57 -04:00
|
|
|
) -> List[Tuple[Any, ...]]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
Executes a SELECT query on the named table with start and limit,
|
|
|
|
of row numbers, which may return zero or number of rows from start to limit,
|
|
|
|
returning the result as a list of dicts.
|
|
|
|
|
2019-12-06 06:28:44 -05:00
|
|
|
Use `filters` to search attributes using SQL wildcards and/or `keyvalues` to
|
|
|
|
select attributes with exact matches. All constraints are joined together
|
|
|
|
using 'AND'.
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
Args:
|
2020-08-20 09:00:59 -04:00
|
|
|
txn: Transaction object
|
|
|
|
table: the table name
|
|
|
|
orderby: Column to order the results by.
|
|
|
|
start: Index to begin the query at.
|
|
|
|
limit: Number of results to return.
|
|
|
|
retcols: the names of the columns to return
|
|
|
|
filters:
|
2019-12-06 06:28:44 -05:00
|
|
|
column names and values to filter the rows with, or None to not
|
|
|
|
apply a WHERE ? LIKE ? clause.
|
2020-08-20 09:00:59 -04:00
|
|
|
keyvalues:
|
2019-12-06 06:28:44 -05:00
|
|
|
column names and values to select the rows with, or None to not
|
2021-03-25 06:34:23 -04:00
|
|
|
apply a WHERE key = value clause.
|
|
|
|
exclude_keyvalues:
|
|
|
|
column names and values to exclude rows with, or None to not
|
|
|
|
apply a WHERE key != value clause.
|
2020-08-20 09:00:59 -04:00
|
|
|
order_direction: Whether the results should be ordered "ASC" or "DESC".
|
|
|
|
|
2019-12-04 08:52:46 -05:00
|
|
|
Returns:
|
2023-10-06 11:41:57 -04:00
|
|
|
The result as a list of tuples.
|
2019-12-04 08:52:46 -05:00
|
|
|
"""
|
|
|
|
if order_direction not in ["ASC", "DESC"]:
|
|
|
|
raise ValueError("order_direction must be one of 'ASC' or 'DESC'.")
|
|
|
|
|
2021-03-25 06:34:23 -04:00
|
|
|
where_clause = "WHERE " if filters or keyvalues or exclude_keyvalues else ""
|
2021-07-15 12:46:54 -04:00
|
|
|
arg_list: List[Any] = []
|
2019-12-06 06:28:44 -05:00
|
|
|
if filters:
|
|
|
|
where_clause += " AND ".join("%s LIKE ?" % (k,) for k in filters)
|
|
|
|
arg_list += list(filters.values())
|
|
|
|
where_clause += " AND " if filters and keyvalues else ""
|
2019-12-04 08:52:46 -05:00
|
|
|
if keyvalues:
|
2019-12-06 06:28:44 -05:00
|
|
|
where_clause += " AND ".join("%s = ?" % (k,) for k in keyvalues)
|
|
|
|
arg_list += list(keyvalues.values())
|
2021-03-25 06:34:23 -04:00
|
|
|
if exclude_keyvalues:
|
|
|
|
where_clause += " AND ".join("%s != ?" % (k,) for k in exclude_keyvalues)
|
|
|
|
arg_list += list(exclude_keyvalues.values())
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
sql = "SELECT %s FROM %s %s ORDER BY %s %s LIMIT ? OFFSET ?" % (
|
|
|
|
", ".join(retcols),
|
|
|
|
table,
|
|
|
|
where_clause,
|
|
|
|
orderby,
|
|
|
|
order_direction,
|
|
|
|
)
|
2019-12-06 06:28:44 -05:00
|
|
|
txn.execute(sql, arg_list + [limit, start])
|
2019-12-04 08:52:46 -05:00
|
|
|
|
2023-10-06 11:41:57 -04:00
|
|
|
return txn.fetchall()
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
|
|
|
|
def make_in_list_sql_clause(
|
2024-05-29 08:19:10 -04:00
|
|
|
database_engine: BaseDatabaseEngine,
|
|
|
|
column: str,
|
|
|
|
iterable: Collection[Any],
|
|
|
|
*,
|
|
|
|
negative: bool = False,
|
2020-02-19 10:47:11 -05:00
|
|
|
) -> Tuple[str, list]:
|
2019-12-04 08:52:46 -05:00
|
|
|
"""Returns an SQL clause that checks the given column is in the iterable.
|
|
|
|
|
|
|
|
On SQLite this expands to `column IN (?, ?, ...)`, whereas on Postgres
|
|
|
|
it expands to `column = ANY(?)`. While both DBs support the `IN` form,
|
|
|
|
using the `ANY` form on postgres means that it views queries with
|
|
|
|
different length iterables as the same, helping the query stats.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
database_engine
|
|
|
|
column: Name of the column
|
|
|
|
iterable: The values to check the column against.
|
2024-05-29 08:19:10 -04:00
|
|
|
negative: Whether we should check for inequality, i.e. `NOT IN`
|
2019-12-04 08:52:46 -05:00
|
|
|
|
|
|
|
Returns:
|
|
|
|
A tuple of SQL query and the args
|
|
|
|
"""
|
|
|
|
|
|
|
|
if database_engine.supports_using_any_list:
|
|
|
|
# This should hopefully be faster, but also makes postgres query
|
|
|
|
# stats easier to understand.
|
2024-05-29 08:19:10 -04:00
|
|
|
if not negative:
|
|
|
|
clause = f"{column} = ANY(?)"
|
|
|
|
else:
|
|
|
|
clause = f"{column} != ALL(?)"
|
|
|
|
|
|
|
|
return clause, [list(iterable)]
|
2019-12-04 08:52:46 -05:00
|
|
|
else:
|
2024-05-29 08:19:10 -04:00
|
|
|
params = ",".join("?" for _ in iterable)
|
|
|
|
if not negative:
|
|
|
|
clause = f"{column} IN ({params})"
|
|
|
|
else:
|
|
|
|
clause = f"{column} NOT IN ({params})"
|
|
|
|
return clause, list(iterable)
|
2020-04-07 18:06:39 -04:00
|
|
|
|
|
|
|
|
2022-10-03 08:46:36 -04:00
|
|
|
# These overloads ensure that `columns` and `iterable` values have the same length.
|
|
|
|
# Suppress "Single overload definition, multiple required" complaint.
|
|
|
|
@overload # type: ignore[misc]
|
|
|
|
def make_tuple_in_list_sql_clause(
|
|
|
|
database_engine: BaseDatabaseEngine,
|
|
|
|
columns: Tuple[str, str],
|
|
|
|
iterable: Collection[Tuple[Any, Any]],
|
2024-03-13 12:46:44 -04:00
|
|
|
) -> Tuple[str, list]: ...
|
2022-10-03 08:46:36 -04:00
|
|
|
|
|
|
|
|
|
|
|
def make_tuple_in_list_sql_clause(
|
|
|
|
database_engine: BaseDatabaseEngine,
|
|
|
|
columns: Tuple[str, ...],
|
|
|
|
iterable: Collection[Tuple[Any, ...]],
|
|
|
|
) -> Tuple[str, list]:
|
|
|
|
"""Returns an SQL clause that checks the given tuple of columns is in the iterable.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
database_engine
|
|
|
|
columns: Names of the columns in the tuple.
|
|
|
|
iterable: The tuples to check the columns against.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
A tuple of SQL query and the args
|
|
|
|
"""
|
|
|
|
if len(columns) == 0:
|
|
|
|
# Should be unreachable due to mypy, as long as the overloads are set up right.
|
|
|
|
if () in iterable:
|
|
|
|
return "TRUE", []
|
|
|
|
else:
|
|
|
|
return "FALSE", []
|
|
|
|
|
|
|
|
if len(columns) == 1:
|
|
|
|
# Use `= ANY(?)` on postgres.
|
|
|
|
return make_in_list_sql_clause(
|
|
|
|
database_engine, next(iter(columns)), [values[0] for values in iterable]
|
|
|
|
)
|
|
|
|
|
|
|
|
# There are multiple columns. Avoid using an `= ANY(?)` clause on postgres, as
|
|
|
|
# indices are not used when there are multiple columns. Instead, use an `IN`
|
|
|
|
# expression.
|
|
|
|
#
|
|
|
|
# `IN ((?, ...), ...)` with tuples is supported by postgres only, whereas
|
|
|
|
# `IN (VALUES (?, ...), ...)` is supported by both sqlite and postgres.
|
|
|
|
# Thus, the latter is chosen.
|
|
|
|
|
|
|
|
if len(iterable) == 0:
|
|
|
|
# A 0-length `VALUES` list is not allowed in sqlite or postgres.
|
|
|
|
# Also note that a 0-length `IN (...)` clause (not using `VALUES`) is not
|
|
|
|
# allowed in postgres.
|
|
|
|
return "FALSE", []
|
|
|
|
|
|
|
|
tuple_sql = "(%s)" % (",".join("?" for _ in columns),)
|
|
|
|
return "(%s) IN (VALUES %s)" % (
|
|
|
|
",".join(column for column in columns),
|
|
|
|
",".join(tuple_sql for _ in iterable),
|
|
|
|
), [value for values in iterable for value in values]
|
|
|
|
|
|
|
|
|
2020-04-07 18:06:39 -04:00
|
|
|
KV = TypeVar("KV")
|
|
|
|
|
|
|
|
|
2021-04-08 13:29:57 -04:00
|
|
|
def make_tuple_comparison_clause(keys: List[Tuple[str, KV]]) -> Tuple[str, List[KV]]:
|
2020-04-07 18:06:39 -04:00
|
|
|
"""Returns a tuple comparison SQL clause
|
|
|
|
|
2021-04-08 08:45:19 -04:00
|
|
|
Builds a SQL clause that looks like "(a, b) > (?, ?)"
|
2020-04-07 18:06:39 -04:00
|
|
|
|
|
|
|
Args:
|
|
|
|
keys: A set of (column, value) pairs to be compared.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
A tuple of SQL query and the args
|
|
|
|
"""
|
2021-04-08 08:45:19 -04:00
|
|
|
return (
|
|
|
|
"(%s) > (%s)" % (",".join(k[0] for k in keys), ",".join("?" for _ in keys)),
|
|
|
|
[k[1] for k in keys],
|
|
|
|
)
|