Fix paginating /relations with a live token (#14866)

The `/relations` endpoint was not properly handle "live tokens"
(i.e sync tokens), to do this properly we abstract the code that
`/messages` has and re-use it.
This commit is contained in:
Patrick Cloke 2023-01-26 13:24:15 -05:00 committed by GitHub
parent ba79fb4a61
commit 345576bc34
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 123 additions and 70 deletions

View file

@ -170,6 +170,104 @@ def generate_pagination_where_clause(
return " AND ".join(where_clause)
def generate_pagination_bounds(
direction: str,
from_token: Optional[RoomStreamToken],
to_token: Optional[RoomStreamToken],
) -> Tuple[
str, Optional[Tuple[Optional[int], int]], Optional[Tuple[Optional[int], int]]
]:
"""
Generate a start and end point for this page of events.
Args:
direction: Whether pagination is going forwards or backwards. One of "f" or "b".
from_token: The token to start pagination at, or None to start at the first value.
to_token: The token to end pagination at, or None to not limit the end point.
Returns:
A three tuple of:
ASC or DESC for sorting of the query.
The starting position as a tuple of ints representing
(topological position, stream position) or None if no from_token was
provided. The topological position may be None for live tokens.
The end position in the same format as the starting position, or None
if no to_token was provided.
"""
# Tokens really represent positions between elements, but we use
# the convention of pointing to the event before the gap. Hence
# we have a bit of asymmetry when it comes to equalities.
if direction == "b":
order = "DESC"
else:
order = "ASC"
# The bounds for the stream tokens are complicated by the fact
# that we need to handle the instance_map part of the tokens. We do this
# by fetching all events between the min stream token and the maximum
# stream token (as returned by `RoomStreamToken.get_max_stream_pos`) and
# then filtering the results.
from_bound: Optional[Tuple[Optional[int], int]] = None
if from_token:
if from_token.topological is not None:
from_bound = from_token.as_historical_tuple()
elif direction == "b":
from_bound = (
None,
from_token.get_max_stream_pos(),
)
else:
from_bound = (
None,
from_token.stream,
)
to_bound: Optional[Tuple[Optional[int], int]] = None
if to_token:
if to_token.topological is not None:
to_bound = to_token.as_historical_tuple()
elif direction == "b":
to_bound = (
None,
to_token.stream,
)
else:
to_bound = (
None,
to_token.get_max_stream_pos(),
)
return order, from_bound, to_bound
def generate_next_token(
direction: str, last_topo_ordering: int, last_stream_ordering: int
) -> RoomStreamToken:
"""
Generate the next room stream token based on the currently returned data.
Args:
direction: Whether pagination is going forwards or backwards. One of "f" or "b".
last_topo_ordering: The last topological ordering being returned.
last_stream_ordering: The last stream ordering being returned.
Returns:
A new RoomStreamToken to return to the client.
"""
if direction == "b":
# Tokens are positions between events.
# This token points *after* the last event in the chunk.
# We need it to point to the event before it in the chunk
# when we are going backwards so we subtract one from the
# stream part.
last_stream_ordering -= 1
return RoomStreamToken(last_topo_ordering, last_stream_ordering)
def _make_generic_sql_bound(
bound: str,
column_names: Tuple[str, str],
@ -1300,47 +1398,11 @@ class StreamWorkerStore(EventsWorkerStore, SQLBaseStore):
`to_token`), or `limit` is zero.
"""
# Tokens really represent positions between elements, but we use
# the convention of pointing to the event before the gap. Hence
# we have a bit of asymmetry when it comes to equalities.
args = [False, room_id]
if direction == "b":
order = "DESC"
else:
order = "ASC"
# The bounds for the stream tokens are complicated by the fact
# that we need to handle the instance_map part of the tokens. We do this
# by fetching all events between the min stream token and the maximum
# stream token (as returned by `RoomStreamToken.get_max_stream_pos`) and
# then filtering the results.
if from_token.topological is not None:
from_bound: Tuple[Optional[int], int] = from_token.as_historical_tuple()
elif direction == "b":
from_bound = (
None,
from_token.get_max_stream_pos(),
)
else:
from_bound = (
None,
from_token.stream,
)
to_bound: Optional[Tuple[Optional[int], int]] = None
if to_token:
if to_token.topological is not None:
to_bound = to_token.as_historical_tuple()
elif direction == "b":
to_bound = (
None,
to_token.stream,
)
else:
to_bound = (
None,
to_token.get_max_stream_pos(),
)
order, from_bound, to_bound = generate_pagination_bounds(
direction, from_token, to_token
)
bounds = generate_pagination_where_clause(
direction=direction,
@ -1436,16 +1498,10 @@ class StreamWorkerStore(EventsWorkerStore, SQLBaseStore):
][:limit]
if rows:
topo = rows[-1].topological_ordering
token = rows[-1].stream_ordering
if direction == "b":
# Tokens are positions between events.
# This token points *after* the last event in the chunk.
# We need it to point to the event before it in the chunk
# when we are going backwards so we subtract one from the
# stream part.
token -= 1
next_token = RoomStreamToken(topo, token)
assert rows[-1].topological_ordering is not None
next_token = generate_next_token(
direction, rows[-1].topological_ordering, rows[-1].stream_ordering
)
else:
# TODO (erikj): We should work out what to do here instead.
next_token = to_token if to_token else from_token