Limit size of get_auth_chain_ids query (#6947)

This commit is contained in:
Erik Johnston 2020-02-19 09:39:27 +00:00 committed by GitHub
parent adfaea8c69
commit 5a5abd55e8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 23 additions and 17 deletions

1
changelog.d/6947.misc Normal file
View File

@ -0,0 +1 @@
Increase perf of `get_auth_chain_ids` used in state res v2.

View File

@ -62,32 +62,37 @@ class EventFederationWorkerStore(EventsWorkerStore, SignatureWorkerStore, SQLBas
)
def _get_auth_chain_ids_txn(self, txn, event_ids, include_given):
if include_given:
results = set(event_ids)
else:
results = set()
if isinstance(self.database_engine, PostgresEngine):
# For efficiency we make the database do this if we can.
sql = """
WITH RECURSIVE auth_chain(event_id) AS (
SELECT auth_id FROM event_auth WHERE event_id = ANY(?)
UNION
SELECT auth_id FROM event_auth
INNER JOIN auth_chain USING (event_id)
)
SELECT event_id FROM auth_chain
"""
txn.execute(sql, (list(event_ids),))
results = set(event_id for event_id, in txn)
# We need to be a little careful with querying large amounts at
# once, for some reason postgres really doesn't like it. We do this
# by only asking for auth chain of 500 events at a time.
event_ids = list(event_ids)
chunks = [event_ids[x : x + 500] for x in range(0, len(event_ids), 500)]
for chunk in chunks:
sql = """
WITH RECURSIVE auth_chain(event_id) AS (
SELECT auth_id FROM event_auth WHERE event_id = ANY(?)
UNION
SELECT auth_id FROM event_auth
INNER JOIN auth_chain USING (event_id)
)
SELECT event_id FROM auth_chain
"""
txn.execute(sql, (chunk,))
if include_given:
results.update(event_ids)
results.update(event_id for event_id, in txn)
return list(results)
# Database doesn't necessarily support recursive CTE, so we fall
# back to do doing it manually.
if include_given:
results = set(event_ids)
else:
results = set()
base_sql = "SELECT auth_id FROM event_auth WHERE "