arminius_lib_db_account_data/

lib.rs

// SPDX-FileCopyrightText: 2025 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0

//! Queries relating to account data
//!
//!
//! # Example for tags
//!
//! Content of a `room_account_data` row to store room tags:
//!
//!
//! | user_id  | room_id                               | data_type   |                       data                          |
//! -----------|---------------------------------------|-------------|-----------------------------------------------------|
//! | 1        | `!kZrfCE0UDTT3PjUosE:localhost:44199` | `m.tag`     | `{"tags": {"u.valid.ordered.tag": {"order": 0.3}}}` |
//!
//!
//! A Matrix client requesting its tag would receive:
//!
//! ```json
//! {
//!     "tags": {
//!         "u.valid.ordered.tag": {
//!             "order": 0.3
//!         }
//!     }
//! }
//! ```
//!
//! Which corresponds to the following `m.tag` Matrix event:
//!
//! ```json
//! {
//!     "content": {
//!         "tags": {
//!             "u.valid.ordered.tag": {
//!                 "order": 0.3
//!             }
//!         }
//!     },
//!     "type": "m.tag"
//! }
//! ```

use arminius_lib_db::{Database, DatabaseError, ResultExt as _, TransactedDb};
use arminius_lib_model::types::StreamId;
use ruma::{
	OwnedRoomId, OwnedUserId, RoomId, UserId,
	api::client::filter::{Filter, RoomEventFilter},
	events::{
		GlobalAccountDataEventType, RoomAccountDataEvent, RoomAccountDataEventContent,
		RoomAccountDataEventType, room::member::MembershipState,
	},
};
use serde::de::DeserializeOwned;
use serde_json::json;
use sqlx::postgres::PgQueryResult;

/// Used to query room account data events along their room ID.
#[derive(Debug)]
pub struct SyncRoomAccountDataEvent {
	/// The room ID the event belongs to
	pub room_id: OwnedRoomId,
	/// The user's membership in the event's room, at the moment it was fetched.
	pub membership: MembershipState,
	/// The stream ID of that event.
	pub stream_id: StreamId,
	/// The actual event as JSON
	pub event: serde_json::Value,
}

/// Used to query global account data events, along their stream ID.
#[derive(Debug, PartialEq)]
pub struct SyncGlobalAccountDataEvent {
	/// The stream ID of that event.
	pub stream_id: StreamId,
	/// The actual event as JSON
	pub event: serde_json::Value,
}

/// Extension trait to acquire a per-user advisory lock for account data
/// writes, held for the duration of the current transaction.
pub trait AccountDataLock: AccountDataQueries {
	/// Acquire a transaction-scoped, per-user advisory lock for account data
	/// writes.
	///
	/// This serializes every writer performing a read-modify-write on a
	/// user's account data (e.g. the tag handlers, which read the `m.tag`
	/// blob, modify it in Rust, and upsert it back). It must be called
	/// *before* the read that a subsequent write depends on, so that the
	/// whole read-then-write sequence is atomic with respect to other
	/// writers. Without it, two concurrent writers can both read the same
	/// blob and the loser's update is silently overwritten (lost update,
	/// see issue #688).
	///
	/// `pg_advisory_xact_lock` is preferred over `SELECT ... FOR UPDATE` on
	/// the `room_account_data` row: `FOR UPDATE` only locks rows that already
	/// exist, so two concurrent writers inserting the *first* row for a
	/// `(user, room, type)` key would still race. The advisory lock also
	/// releases automatically on commit or rollback and is re-entrant within
	/// the same session.
	///
	/// The two-key `(room_account_data_oid, users.nid)` form is deliberate —
	/// do NOT reduce it to `pg_advisory_xact_lock(nid::bigint)`. A small
	/// `nid` in the `bigint` form is stored as `classid = 0`, the space
	/// `max_safe_seq` scans for in-flight stream id values (see
	/// `00_locking_sequence.sql`). Since account data writes allocate a
	/// stream id in the same transaction, a `classid = 0` lock would collapse
	/// the sync watermark to the user's `nid` and stall /sync. Namespacing by
	/// table OID keeps the lock out of that scan.
	///
	/// Locking on `users.nid` serializes all account data writes of one user
	/// (across rooms and data types). That is coarser than strictly needed,
	/// but account data writes are low-frequency and this keeps the lock key
	/// a plain `int4` without hashing.
	///
	/// If the user does not exist no lock is taken — handlers only run for
	/// authenticated (i.e. existing) users.
	#[tracing::instrument(skip(self))]
	async fn lock_room_account_data(&mut self, user_id: &UserId) -> Result<(), DatabaseError> {
		sqlx::query_scalar!(
			"SELECT pg_advisory_xact_lock(
			    (SELECT oid::int FROM pg_class WHERE relname = 'room_account_data' AND relkind = 'r'),
			    nid
			)
			FROM users WHERE user_id = $1",
			user_id.as_str(),
		)
		.fetch_optional(self.connection().await?)
		.await?;

		Ok(())
	}

	/// Retrieve room account data for a user in a specific room, taking the
	/// per-user account data lock ([`Self::lock_room_account_data`]) first.
	///
	/// Use this instead of `fetch_room_account_data` whenever the fetched
	/// value is modified and written back within the same transaction
	/// (read-modify-write), so concurrent writers are serialized and cannot
	/// overwrite each other's updates (#688).
	#[tracing::instrument(skip(self))]
	async fn fetch_room_account_data_for_update<C>(
		&mut self,
		user_id: &UserId,
		room_id: &RoomId,
		event_type: &RoomAccountDataEventType,
	) -> Result<Option<RoomAccountDataEvent<C>>, DatabaseError>
	where
		C: RoomAccountDataEventContent + DeserializeOwned,
		Self: Sized,
	{
		self.lock_room_account_data(user_id).await?;

		self.fetch_room_account_data(user_id, room_id, event_type).await
	}
}

/// The lock is transaction-scoped (`pg_advisory_xact_lock` releases on commit
/// or rollback), so it is only implemented for [`TransactedDb`]: taking it on
/// a plain connection would release it again right after the statement,
/// providing no protection.
impl AccountDataLock for TransactedDb<'_> {}

/// Queries relating to account data
// Queries are usually used as extension traits so suggest `import as _`.
#[rust_analyzer::prefer_underscore_import]
pub trait AccountDataQueries: Database {
	/// Inserts or updates room account data for a user in a specific room.
	#[tracing::instrument(skip(self))]
	async fn upsert_room_account_data(
		&mut self,
		user_id: &UserId,
		room_id: &RoomId,
		event_type: &RoomAccountDataEventType,
		data: serde_json::Value,
	) -> Result<PgQueryResult, DatabaseError> {
		sqlx::query!(
			"
			INSERT INTO room_account_data (user_nid, room_nid, data_type, data)
				VALUES (
					(SELECT nid FROM users WHERE user_id = $1),
					(SELECT nid FROM rooms WHERE room_id = $2),
					$3,
					$4
				)
			ON CONFLICT (user_nid, room_nid, data_type)
				DO UPDATE SET
					data = EXCLUDED.data,
					stream_id = EXCLUDED.stream_id
			",
			user_id.as_str(),
			room_id.as_str(),
			event_type.to_string(),
			data,
		)
		.execute(self.connection().await?)
		.await
		.map_err(Into::into)
	}

	/// Inserts or updates global account data for a user.
	#[tracing::instrument(skip(self))]
	async fn upsert_global_account_data(
		&mut self,
		user_id: &UserId,
		event_type: &GlobalAccountDataEventType,
		data: serde_json::Value,
	) -> Result<PgQueryResult, DatabaseError> {
		sqlx::query!(
			"
			 INSERT INTO global_account_data (user_nid, data_type, data)
			 	VALUES (
			 		(SELECT nid FROM users WHERE user_id = $1),
			 		$2,
			 		$3
			 	)
			 ON CONFLICT (user_nid, data_type)
			 	DO UPDATE SET
			 		data = EXCLUDED.data,
					stream_id = EXCLUDED.stream_id
			",
			user_id.as_str(),
			event_type.to_string(),
			data,
		)
		.execute(self.connection().await?)
		.await
		.map_err(Into::into)
	}

	/// Retrieve room account data for a user in a specific room.
	#[tracing::instrument(skip(self))]
	async fn fetch_room_account_data<C: RoomAccountDataEventContent + DeserializeOwned>(
		&mut self,
		user_id: &UserId,
		room_id: &RoomId,
		event_type: &RoomAccountDataEventType,
	) -> Result<Option<RoomAccountDataEvent<C>>, DatabaseError> {
		self.fetch_room_account_data_untyped(user_id, room_id, event_type)
			.await?
			.map(|data| {
				let event_content = serde_json::from_value::<C>(data)
					.invalid_data("Invalid room account data event in database")?;

				Ok(RoomAccountDataEvent::new(event_content))
			})
			.transpose()
	}

	/// Retrieve room account data for a user in a specific room.
	///
	/// This function, as opposed to [`Self::fetch_room_account_data`], is to be
	/// used when you do not need to know which variant of
	/// [`RoomAccountDataEvent`] you have in the caller.
	#[tracing::instrument(skip(self))]
	async fn fetch_room_account_data_untyped(
		&mut self,
		user_id: &UserId,
		room_id: &RoomId,
		event_type: &RoomAccountDataEventType,
	) -> Result<Option<serde_json::Value>, DatabaseError> {
		Ok(sqlx::query!(
			"
			SELECT
				room_account_data.data
			FROM
				room_account_data
				JOIN users ON room_account_data.user_nid = users.nid
				JOIN rooms ON room_account_data.room_nid = rooms.nid
			WHERE
				users.user_id = $1
				AND room_id = $2
				AND data_type = $3
			",
			user_id.as_str(),
			room_id.as_str(),
			event_type.to_string(),
		)
		.fetch_optional(self.connection().await?)
		.await?
		.map(|row| row.data))
	}

	/// Imports a global account-data row, preserving the source stream_id.
	/// Idempotent: an existing (user, data_type) row is left untouched.
	///
	/// Preserving the stream_id bypasses `account_data_streamids`, so the
	/// sequence is advanced past the imported value in the same transaction.
	/// This keeps the "next auto-assigned stream_id never collides with an
	/// existing row" invariant: the import leaves the database healthy on its
	/// own, with no external sequence reset to remember.
	#[tracing::instrument(skip(self))]
	async fn import_global_account_data(
		&mut self,
		user_id: &UserId,
		data_type: &str,
		data: serde_json::Value,
		stream_id: i64,
	) -> Result<bool, DatabaseError> {
		let result = sqlx::query!(
			r#"
			INSERT INTO global_account_data (user_nid, data_type, data, stream_id)
				VALUES ((
						SELECT
							nid
						FROM
							users
						WHERE
							user_id = $1), $2, $3, $4)
			ON CONFLICT (user_nid,
				data_type)
				DO NOTHING
			"#,
			user_id.as_str(),
			data_type,
			data,
			stream_id,
		)
		.execute(self.connection().await?)
		.await?;
		advance_account_data_sequence(self, stream_id).await?;
		Ok(result.rows_affected() > 0)
	}

	/// Imports a room account-data row, preserving the source stream_id.
	/// Idempotent: an existing (user, room, data_type) row is left untouched.
	///
	/// Shares `account_data_streamids` with
	/// [`Self::import_global_account_data`] and, like it, advances the
	/// sequence past the preserved stream_id in the same transaction so the
	/// import leaves the database healthy on its own.
	#[tracing::instrument(skip(self))]
	async fn import_room_account_data(
		&mut self,
		user_id: &UserId,
		room_id: &RoomId,
		data_type: &str,
		data: serde_json::Value,
		stream_id: i64,
	) -> Result<bool, DatabaseError> {
		let result = sqlx::query!(
			r#"
			INSERT INTO room_account_data (user_nid, room_nid, data_type, data, stream_id)
			VALUES
				((
						SELECT
							nid
						FROM
							users
						WHERE
							user_id = $1), (
							SELECT
								nid
							FROM
								rooms
							WHERE
								room_id = $2), $3, $4, $5)
				ON CONFLICT (user_nid,
					room_nid,
					data_type)
				DO NOTHING
			"#,
			user_id.as_str(),
			room_id.as_str(),
			data_type,
			data,
			stream_id,
		)
		.execute(self.connection().await?)
		.await?;
		advance_account_data_sequence(self, stream_id).await?;
		Ok(result.rows_affected() > 0)
	}

	/// Fetch global account data for a user based on the event type.
	#[tracing::instrument(skip(self))]
	async fn fetch_global_account_data_untyped(
		&mut self,
		user_id: &UserId,
		event_type: &GlobalAccountDataEventType,
	) -> Result<Option<serde_json::Value>, DatabaseError> {
		Ok(sqlx::query!(
			"
			SELECT
				global_account_data.data
			FROM
				global_account_data
				JOIN users ON global_account_data.user_nid = users.nid
			WHERE
				users.user_id = $1
				AND data_type = $2
			",
			user_id.as_str(),
			event_type.to_string(),
		)
		.fetch_optional(self.connection().await?)
		.await?
		.map(|row| row.data))
	}

	/// Fetches any entry in the `global_account_data` table which `stream_id`
	/// strictly higher than `from_stream_id`.
	#[tracing::instrument(skip(self))]
	async fn fetch_sync_global_account_data_events(
		&mut self,
		user_id: &UserId,
		from_stream_id: StreamId,
		filter: Option<&Filter>,
	) -> Result<Vec<SyncGlobalAccountDataEvent>, DatabaseError> {
		let types = filter.and_then(|filter| filter.types.clone());
		let not_types = filter.map(|filter| filter.not_types.clone()).unwrap_or_default();
		Ok(sqlx::query!(
			"
			SELECT
				stream_id,
				data_type,
				data
			FROM
				global_account_data
				JOIN users ON global_account_data.user_nid = users.nid
			WHERE
				users.user_id = $1
				AND stream_id > $2
				AND ($3::text[] IS NULL OR data_type::text = ANY ($3::text[]))
				AND NOT (data_type = ANY ($4::text[]))
				AND stream_id <= max_safe_seq ('account_data_streamids')
			",
			user_id.to_string(),
			from_stream_id,
			types.as_deref(),
			&not_types
		)
		.fetch_all(self.connection().await?)
		.await?
		.into_iter()
		.map(|row| SyncGlobalAccountDataEvent {
			stream_id: row.stream_id,
			event: json!({
				"type": row.data_type,
				"content": row.data
			}),
		})
		.collect())
	}

	/// Fetches any entry in the `room_account_data` table which `stream_id`
	/// strictly higher than `from_stream_id`.
	#[tracing::instrument(skip(self))]
	async fn fetch_sync_room_account_data_events(
		&mut self,
		user_id: &OwnedUserId,
		from_stream_id: StreamId,
		filter: Option<&RoomEventFilter>,
	) -> Result<Vec<SyncRoomAccountDataEvent>, DatabaseError> {
		let types = filter.and_then(|filter| filter.types.clone());
		let not_types = filter.map(|filter| filter.not_types.clone()).unwrap_or_default();
		sqlx::query!(
			r#"
			SELECT
				room_membership_state.room_id AS "room_id!",
				room_membership_state.membership AS "membership!",
				stream_id,
				data_type,
				data
			FROM
				room_account_data
				JOIN room_membership_state ($1) ON room_account_data.room_nid = room_membership_state.room_nid
			WHERE
				room_membership_state.membership IN ('join', 'leave')
				AND stream_id > $2
				AND ($3::text[] IS NULL
					OR data_type::text = ANY ($3::text[]))
				AND NOT (data_type = ANY ($4::text[]))
				AND stream_id <= max_safe_seq ('account_data_streamids')
			"#,
			user_id.to_string(),
			from_stream_id,
			types.as_deref(),
			&not_types
		)
		.fetch_all(self.connection().await?)
		.await?
		.into_iter()
		.map(|row| {
			let room_id = RoomId::parse(row.room_id).invalid_data("Invalid room ID in database")?;

			Ok(SyncRoomAccountDataEvent {
				room_id,
				membership: MembershipState::from(row.membership),
				stream_id: row.stream_id,
				event: json!({
					"type": row.data_type,
					"content": row.data
				}),
			})
		})
		.collect::<Result<Vec<_>, _>>()
	}
}

impl<T> AccountDataQueries for T where T: Database {}

/// Advances `account_data_streamids` so its last value covers `stream_id`,
/// never moving it backwards. Called by the account-data importers after they
/// insert a row carrying a preserved (sequence-bypassing) stream_id, so the
/// next `nextval` cannot hand out an id that already exists.
///
/// `setval` is non-transactional, so on a rolled-back import batch the sequence
/// may stay advanced, that only ever skips ids forward, which is harmless
/// (stream_ids need only be monotonic, not gap-free).
async fn advance_account_data_sequence<T: Database + ?Sized>(
	db: &mut T,
	stream_id: i64,
) -> Result<(), DatabaseError> {
	sqlx::query_scalar!(
		"SELECT setval(
			'account_data_streamids',
			GREATEST((SELECT last_value FROM account_data_streamids), $1),
			true)",
		stream_id,
	)
	.fetch_one(db.connection().await?)
	.await?;
	Ok(())
}