Move DB pool and helper functions into dedicated Database class

This commit is contained in:
Erik Johnston 2019-12-04 13:52:46 +00:00
parent ddbbfc9512
commit 756d4942f5
62 changed files with 2377 additions and 2295 deletions

View file

@ -45,13 +45,13 @@ class StatsRoomTests(unittest.HomeserverTestCase):
self.store._all_done = False
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{"update_name": "populate_stats_prepare", "progress_json": "{}"},
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_rooms",
@ -61,7 +61,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_users",
@ -71,7 +71,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_cleanup",
@ -82,7 +82,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
def get_all_room_state(self):
return self.store.simple_select_list(
return self.store.db.simple_select_list(
"room_stats_state", None, retcols=("name", "topic", "canonical_alias")
)
@ -96,7 +96,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
end_ts = self.store.quantise_stats_time(self.reactor.seconds() * 1000)
return self.get_success(
self.store.simple_select_one(
self.store.db.simple_select_one(
table + "_historical",
{id_col: stat_id, end_ts: end_ts},
cols,
@ -180,7 +180,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
self.handler.stats_enabled = True
self.store._all_done = False
self.get_success(
self.store.simple_update_one(
self.store.db.simple_update_one(
table="stats_incremental_position",
keyvalues={},
updatevalues={"stream_id": 0},
@ -188,7 +188,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{"update_name": "populate_stats_prepare", "progress_json": "{}"},
)
@ -205,13 +205,13 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# Now do the initial ingestion.
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{"update_name": "populate_stats_process_rooms", "progress_json": "{}"},
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_cleanup",
@ -656,12 +656,12 @@ class StatsRoomTests(unittest.HomeserverTestCase):
self.store._all_done = False
self.get_success(
self.store.simple_delete(
self.store.db.simple_delete(
"room_stats_current", {"1": 1}, "test_delete_stats"
)
)
self.get_success(
self.store.simple_delete(
self.store.db.simple_delete(
"user_stats_current", {"1": 1}, "test_delete_stats"
)
)
@ -675,7 +675,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
self.store._all_done = False
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_rooms",
@ -685,7 +685,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_users",
@ -695,7 +695,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_cleanup",

View file

@ -158,7 +158,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
def get_users_in_public_rooms(self):
r = self.get_success(
self.store.simple_select_list(
self.store.db.simple_select_list(
"users_in_public_rooms", None, ("user_id", "room_id")
)
)
@ -169,7 +169,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
def get_users_who_share_private_rooms(self):
return self.get_success(
self.store.simple_select_list(
self.store.db.simple_select_list(
"users_who_share_private_rooms",
None,
["user_id", "other_user_id", "room_id"],
@ -184,7 +184,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
self.store._all_done = False
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_user_directory_createtables",
@ -193,7 +193,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_user_directory_process_rooms",
@ -203,7 +203,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_user_directory_process_users",
@ -213,7 +213,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
)
)
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_user_directory_cleanup",

View file

@ -632,7 +632,7 @@ class PurgeRoomTestCase(unittest.HomeserverTestCase):
"state_groups_state",
):
count = self.get_success(
self.store.simple_select_one_onecol(
self.store.db.simple_select_one_onecol(
table=table,
keyvalues={"room_id": room_id},
retcol="COUNT(*)",

View file

@ -323,7 +323,7 @@ class UpsertManyTests(unittest.HomeserverTestCase):
self.table_name = "table_" + hs.get_secrets().token_hex(6)
self.get_success(
self.storage.runInteraction(
self.storage.db.runInteraction(
"create",
lambda x, *a: x.execute(*a),
"CREATE TABLE %s (id INTEGER, username TEXT, value TEXT)"
@ -331,7 +331,7 @@ class UpsertManyTests(unittest.HomeserverTestCase):
)
)
self.get_success(
self.storage.runInteraction(
self.storage.db.runInteraction(
"index",
lambda x, *a: x.execute(*a),
"CREATE UNIQUE INDEX %sindex ON %s(id, username)"
@ -354,9 +354,9 @@ class UpsertManyTests(unittest.HomeserverTestCase):
value_values = [["hello"], ["there"]]
self.get_success(
self.storage.runInteraction(
self.storage.db.runInteraction(
"test",
self.storage.simple_upsert_many_txn,
self.storage.db.simple_upsert_many_txn,
self.table_name,
key_names,
key_values,
@ -367,7 +367,7 @@ class UpsertManyTests(unittest.HomeserverTestCase):
# Check results are what we expect
res = self.get_success(
self.storage.simple_select_list(
self.storage.db.simple_select_list(
self.table_name, None, ["id, username, value"]
)
)
@ -381,9 +381,9 @@ class UpsertManyTests(unittest.HomeserverTestCase):
value_values = [["bleb"]]
self.get_success(
self.storage.runInteraction(
self.storage.db.runInteraction(
"test",
self.storage.simple_upsert_many_txn,
self.storage.db.simple_upsert_many_txn,
self.table_name,
key_names,
key_values,
@ -394,7 +394,7 @@ class UpsertManyTests(unittest.HomeserverTestCase):
# Check results are what we expect
res = self.get_success(
self.storage.simple_select_list(
self.storage.db.simple_select_list(
self.table_name, None, ["id, username, value"]
)
)

View file

@ -37,7 +37,7 @@ class BackgroundUpdateTestCase(unittest.TestCase):
def update(progress, count):
self.clock.advance_time_msec(count * duration_ms)
progress = {"my_key": progress["my_key"] + 1}
yield self.store.runInteraction(
yield self.store.db.runInteraction(
"update_progress",
self.store._background_update_progress_txn,
"test_update",

View file

@ -65,7 +65,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_insert_1col(self):
self.mock_txn.rowcount = 1
yield self.datastore.simple_insert(
yield self.datastore.db.simple_insert(
table="tablename", values={"columname": "Value"}
)
@ -77,7 +77,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_insert_3cols(self):
self.mock_txn.rowcount = 1
yield self.datastore.simple_insert(
yield self.datastore.db.simple_insert(
table="tablename",
# Use OrderedDict() so we can assert on the SQL generated
values=OrderedDict([("colA", 1), ("colB", 2), ("colC", 3)]),
@ -92,7 +92,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
self.mock_txn.rowcount = 1
self.mock_txn.__iter__ = Mock(return_value=iter([("Value",)]))
value = yield self.datastore.simple_select_one_onecol(
value = yield self.datastore.db.simple_select_one_onecol(
table="tablename", keyvalues={"keycol": "TheKey"}, retcol="retcol"
)
@ -106,7 +106,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
self.mock_txn.rowcount = 1
self.mock_txn.fetchone.return_value = (1, 2, 3)
ret = yield self.datastore.simple_select_one(
ret = yield self.datastore.db.simple_select_one(
table="tablename",
keyvalues={"keycol": "TheKey"},
retcols=["colA", "colB", "colC"],
@ -122,7 +122,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
self.mock_txn.rowcount = 0
self.mock_txn.fetchone.return_value = None
ret = yield self.datastore.simple_select_one(
ret = yield self.datastore.db.simple_select_one(
table="tablename",
keyvalues={"keycol": "Not here"},
retcols=["colA"],
@ -137,7 +137,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
self.mock_txn.__iter__ = Mock(return_value=iter([(1,), (2,), (3,)]))
self.mock_txn.description = (("colA", None, None, None, None, None, None),)
ret = yield self.datastore.simple_select_list(
ret = yield self.datastore.db.simple_select_list(
table="tablename", keyvalues={"keycol": "A set"}, retcols=["colA"]
)
@ -150,7 +150,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_update_one_1col(self):
self.mock_txn.rowcount = 1
yield self.datastore.simple_update_one(
yield self.datastore.db.simple_update_one(
table="tablename",
keyvalues={"keycol": "TheKey"},
updatevalues={"columnname": "New Value"},
@ -165,7 +165,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_update_one_4cols(self):
self.mock_txn.rowcount = 1
yield self.datastore.simple_update_one(
yield self.datastore.db.simple_update_one(
table="tablename",
keyvalues=OrderedDict([("colA", 1), ("colB", 2)]),
updatevalues=OrderedDict([("colC", 3), ("colD", 4)]),
@ -180,7 +180,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_delete_one(self):
self.mock_txn.rowcount = 1
yield self.datastore.simple_delete_one(
yield self.datastore.db.simple_delete_one(
table="tablename", keyvalues={"keycol": "Go away"}
)

View file

@ -62,7 +62,9 @@ class CleanupExtremBackgroundUpdateStoreTestCase(HomeserverTestCase):
prepare_database.executescript(txn, schema_path)
self.get_success(
self.store.runInteraction("test_delete_forward_extremities", run_delta_file)
self.store.db.runInteraction(
"test_delete_forward_extremities", run_delta_file
)
)
# Ugh, have to reset this flag

View file

@ -81,7 +81,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
self.pump(0)
result = self.get_success(
self.store.simple_select_list(
self.store.db.simple_select_list(
table="user_ips",
keyvalues={"user_id": user_id},
retcols=["access_token", "ip", "user_agent", "device_id", "last_seen"],
@ -112,7 +112,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
self.pump(0)
result = self.get_success(
self.store.simple_select_list(
self.store.db.simple_select_list(
table="user_ips",
keyvalues={"user_id": user_id},
retcols=["access_token", "ip", "user_agent", "device_id", "last_seen"],
@ -218,7 +218,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
# But clear the associated entry in devices table
self.get_success(
self.store.simple_update(
self.store.db.simple_update(
table="devices",
keyvalues={"user_id": user_id, "device_id": "device_id"},
updatevalues={"last_seen": None, "ip": None, "user_agent": None},
@ -245,7 +245,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
# Register the background update to run again.
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
table="background_updates",
values={
"update_name": "devices_last_seen",
@ -297,7 +297,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
# We should see that in the DB
result = self.get_success(
self.store.simple_select_list(
self.store.db.simple_select_list(
table="user_ips",
keyvalues={"user_id": user_id},
retcols=["access_token", "ip", "user_agent", "device_id", "last_seen"],
@ -323,7 +323,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
# We should get no results.
result = self.get_success(
self.store.simple_select_list(
self.store.db.simple_select_list(
table="user_ips",
keyvalues={"user_id": user_id},
retcols=["access_token", "ip", "user_agent", "device_id", "last_seen"],

View file

@ -61,7 +61,7 @@ class EventFederationWorkerStoreTestCase(tests.unittest.TestCase):
)
for i in range(0, 11):
yield self.store.runInteraction("insert", insert_event, i)
yield self.store.db.runInteraction("insert", insert_event, i)
# this should get the last five and five others
r = yield self.store.get_prev_events_for_room(room_id)
@ -93,9 +93,9 @@ class EventFederationWorkerStoreTestCase(tests.unittest.TestCase):
)
for i in range(0, 20):
yield self.store.runInteraction("insert", insert_event, i, room1)
yield self.store.runInteraction("insert", insert_event, i, room2)
yield self.store.runInteraction("insert", insert_event, i, room3)
yield self.store.db.runInteraction("insert", insert_event, i, room1)
yield self.store.db.runInteraction("insert", insert_event, i, room2)
yield self.store.db.runInteraction("insert", insert_event, i, room3)
# Test simple case
r = yield self.store.get_rooms_with_many_extremities(5, 5, [])

View file

@ -55,7 +55,7 @@ class EventPushActionsStoreTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks
def _assert_counts(noitf_count, highlight_count):
counts = yield self.store.runInteraction(
counts = yield self.store.db.runInteraction(
"", self.store._get_unread_counts_by_pos_txn, room_id, user_id, 0
)
self.assertEquals(
@ -74,7 +74,7 @@ class EventPushActionsStoreTestCase(tests.unittest.TestCase):
yield self.store.add_push_actions_to_staging(
event.event_id, {user_id: action}
)
yield self.store.runInteraction(
yield self.store.db.runInteraction(
"",
self.store._set_push_actions_for_event_and_users_txn,
[(event, None)],
@ -82,12 +82,12 @@ class EventPushActionsStoreTestCase(tests.unittest.TestCase):
)
def _rotate(stream):
return self.store.runInteraction(
return self.store.db.runInteraction(
"", self.store._rotate_notifs_before_txn, stream
)
def _mark_read(stream, depth):
return self.store.runInteraction(
return self.store.db.runInteraction(
"",
self.store._remove_old_push_actions_before_txn,
room_id,
@ -116,7 +116,7 @@ class EventPushActionsStoreTestCase(tests.unittest.TestCase):
yield _inject_actions(6, PlAIN_NOTIF)
yield _rotate(7)
yield self.store.simple_delete(
yield self.store.db.simple_delete(
table="event_push_actions", keyvalues={"1": 1}, desc=""
)
@ -135,7 +135,7 @@ class EventPushActionsStoreTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks
def test_find_first_stream_ordering_after_ts(self):
def add_event(so, ts):
return self.store.simple_insert(
return self.store.db.simple_insert(
"events",
{
"stream_ordering": so,

View file

@ -65,7 +65,7 @@ class MonthlyActiveUsersTestCase(unittest.HomeserverTestCase):
self.store.user_add_threepid(user1, "email", user1_email, now, now)
self.store.user_add_threepid(user2, "email", user2_email, now, now)
self.store.runInteraction(
self.store.db.runInteraction(
"initialise", self.store._initialise_reserved_users, threepids
)
self.pump()
@ -183,7 +183,7 @@ class MonthlyActiveUsersTestCase(unittest.HomeserverTestCase):
)
self.hs.config.mau_limits_reserved_threepids = threepids
self.store.runInteraction(
self.store.db.runInteraction(
"initialise", self.store._initialise_reserved_users, threepids
)
count = self.store.get_monthly_active_count()
@ -244,7 +244,7 @@ class MonthlyActiveUsersTestCase(unittest.HomeserverTestCase):
{"medium": "email", "address": user2_email},
]
self.hs.config.mau_limits_reserved_threepids = threepids
self.store.runInteraction(
self.store.db.runInteraction(
"initialise", self.store._initialise_reserved_users, threepids
)

View file

@ -338,7 +338,7 @@ class RedactionTestCase(unittest.HomeserverTestCase):
)
event_json = self.get_success(
self.store.simple_select_one_onecol(
self.store.db.simple_select_one_onecol(
table="event_json",
keyvalues={"event_id": msg_event.event_id},
retcol="json",
@ -356,7 +356,7 @@ class RedactionTestCase(unittest.HomeserverTestCase):
self.reactor.advance(60 * 60 * 2)
event_json = self.get_success(
self.store.simple_select_one_onecol(
self.store.db.simple_select_one_onecol(
table="event_json",
keyvalues={"event_id": msg_event.event_id},
retcol="json",

View file

@ -132,7 +132,7 @@ class CurrentStateMembershipUpdateTestCase(unittest.HomeserverTestCase):
# Register the background update to run again.
self.get_success(
self.store.simple_insert(
self.store.db.simple_insert(
table="background_updates",
values={
"update_name": "current_state_events_membership",

View file

@ -544,7 +544,7 @@ class HomeserverTestCase(TestCase):
Add the given event as an extremity to the room.
"""
self.get_success(
self.hs.get_datastore().simple_insert(
self.hs.get_datastore().db.simple_insert(
table="event_forward_extremities",
values={"room_id": room_id, "event_id": event_id},
desc="test_add_extremity",