Add some type hints to datastore

[dklimpel]

Pull Request Checklist

• Pull request is based on the develop branch
• Pull request includes a changelog file. The entry should:
  • Be a short description of your change which makes sense to users. "Fixed a bug that prevented receiving messages from other servers." instead of "Moved X method from EventStore to EventWorkerStore.".
  • Use markdown where necessary, mostly for code blocks.
  • End with either a period (.) or an exclamation mark (!).
  • Start with a capital letter.
  • Feel free to credit yourself, by adding a sentence "Contributed by @github_username." or "Contributed by [Your Name]." to the end of the entry.
• Pull request includes a sign off
• Code style is correct
  (run the linters)

Signed-off-by: Dirk Klimpel dirk@klimpel.org

[DMRobertson]

Many thanks Dirk, this looks great. I have a few questions:

[DMRobertson]

This doesn't look right; get_stream_id_for_event_txn passes allow_none=True. Could you explain your thinking?

[dklimpel]

If is set allow_none: bool = False, I am get:

synapse/storage/databases/main/stream.py:786: error: No overload variant of "simple_select_one_onecol_txn" of "DatabasePool" matches argument types "LoggingTransaction", "str", "Dict[str, str]", "str", "bool"  [call-overload]
synapse/storage/databases/main/stream.py:786: note: Possible overload variants:
synapse/storage/databases/main/stream.py:786: note:     def simple_select_one_onecol_txn(cls, txn: LoggingTransaction, table: str, keyvalues: Dict[str, Any], retcol: str, allow_none: Literal[False] = ...) -> Any
synapse/storage/databases/main/stream.py:786: note:     def simple_select_one_onecol_txn(cls, txn: LoggingTransaction, table: str, keyvalues: Dict[str, Any], retcol: str, allow_none: Literal[True] = ...) -> Optional[Any]

I cannot see the mistake:
This should match with:

synapse/synapse/storage/database.py

Lines 1564 to 1571 in 98ec375

	def simple_select_one_onecol_txn(
	cls,
	txn: LoggingTransaction,
	table: str,
	keyvalues: Dict[str, Any],
	retcol: str,
	allow_none: bool = False,
	) -> Optional[Any]:

[DMRobertson]

My understanding:

1. StreamWorkerStore.get_stream_id_for_event_txn is declared to return an int.
2. StreamWorkerStore.get_stream_id returns whatever DatabasePool.simple_select_one_onecol_txn returns. The former passes a value of allow_none through to the latter.
3. DatabasePool.simple_select_one_onecol_txn returns Optional[Any]. If it's called with allow_none=False, then it returns Any; the None case is handled by raising StoreError. (See

   synapse/synapse/storage/database.py

   Lines 1539 to 1561 in e630722

   	@overload
   	@classmethod
   	def simple_select_one_onecol_txn(
   	cls,
   	txn: LoggingTransaction,
   	table: str,
   	keyvalues: Dict[str, Any],
   	retcol: str,
   	allow_none: Literal[False] = False,
   	) -> Any:
   	...
   	@overload
   	@classmethod
   	def simple_select_one_onecol_txn(
   	cls,
   	txn: LoggingTransaction,
   	table: str,
   	keyvalues: Dict[str, Any],
   	retcol: str,
   	allow_none: Literal[True] = True,
   	) -> Optional[Any]:
   	...

   ; this is where the overloads mentioned in the error message come from).

Mypy has spotted that if we call get_stream_id_for_event_txn with allow_none=True, then simple_select_one_onecol_txn might return None. And that would contradict the return type annotation of get_stream_id_for_event_txn.

We only seem to call get_stream_id_for_event_txn in two places, both in this function:

synapse/synapse/storage/databases/main/event_push_actions.py

Lines 184 to 216 in 3e0536c

	def _get_unread_counts_by_receipt_txn(
	self,
	txn: LoggingTransaction,
	room_id: str,
	user_id: str,
	last_read_event_id: Optional[str],
	) -> NotifCounts:
	stream_ordering = None
	if last_read_event_id is not None:
	stream_ordering = self.get_stream_id_for_event_txn( # type: ignore[attr-defined]
	txn,
	last_read_event_id,
	allow_none=True,
	)
	if stream_ordering is None:
	# Either last_read_event_id is None, or it's an event we don't have (e.g.
	# because it's been purged), in which case retrieve the stream ordering for
	# the latest membership event from this user in this room (which we assume is
	# a join).
	event_id = self.db_pool.simple_select_one_onecol_txn(
	txn=txn,
	table="local_current_membership",
	keyvalues={"room_id": room_id, "user_id": user_id},
	retcol="event_id",
	)
	stream_ordering = self.get_stream_id_for_event_txn(txn, event_id) # type: ignore[attr-defined]
	return self._get_unread_counts_by_pos_txn(
	txn, room_id, user_id, stream_ordering
	)

In particular line 194 of the quoted source looks like it's expecting stream_ordering to be possibly None.

[DMRobertson]

Given the above, I think the right thing to do is

1. change get_stream_id_for_event_txn to return Optional[int] instead of int.
2. Try changing EventPushActionsWorkerStore to inherit from StreamWorkerStore and removing the type-ignores where we call get_stream_id_for_event_txn. Doing so might fail due with DataStore MRO problems (related: Improve type hints for data store usage #11165 and more generally our storage classes are hella confusing #11733). If so, that's fine and we can leave things as they are.

Sorry Dirk, that's maybe a bit of information overload. Does that make sense? (If not, is there anything I can clarify?)

[dklimpel]

Does it make sense to open an extra issue for this? And do not touch this type hint here?

[DMRobertson]

That's also fine!

[DMRobertson]

See #12437.

[dklimpel]

Thank you for opening the new issue.

[DMRobertson]

As per your comment, let's leave the annotation of allow_none alone for now; I've raised #12437. Once that's done I think this looks good to merge!

[DMRobertson]

Thanks again, Dirk! The complement test failure should be rectified by #12439.