memory (sqla): Add methods to handle pubsub node subscriptions:
Add `get_pubsub_subscriptions` and `set_pubsub_subscriptions`. rel 476
This commit is contained in:
parent
753623e934
commit
0a0735ec4c
1 changed files with 114 additions and 0 deletions
|
|
@ -81,6 +81,7 @@ from libervia.backend.memory.sqla_mapping import (
|
|||
PubsubItem,
|
||||
PubsubNode,
|
||||
PubsubSub,
|
||||
SubscriptionState,
|
||||
Subject,
|
||||
SyncState,
|
||||
Thread,
|
||||
|
|
@ -1910,6 +1911,119 @@ class Storage:
|
|||
|
||||
await session.commit()
|
||||
|
||||
@aio
|
||||
async def get_pubsub_subscriptions(
|
||||
self,
|
||||
client: SatXMPPEntity,
|
||||
node: PubsubNode | None = None,
|
||||
service: jid.JID | None = None,
|
||||
entity: jid.JID | None = None,
|
||||
node_name: str | None = None,
|
||||
) -> list[PubsubSub]:
|
||||
"""Retrieve subscriptions from the pubsub service.
|
||||
|
||||
@param client: Subscriptions from this client profile will be retrieved.
|
||||
@param node: If set, filter to subscriptions for this specific node instance.
|
||||
Takes precedence over service and node_name. The node will be attached to the
|
||||
results without additional DB query.
|
||||
@param service: If set, filter to subscriptions for this pubsub service.
|
||||
@param entity: If set, filter to subscriptions for this entity (subscriber).
|
||||
@param node_name: If set, filter to subscriptions for this node.
|
||||
@return: PubsubSub instances.
|
||||
@raise exceptions.InternalError: If both node and node_name are provided.
|
||||
"""
|
||||
if node is not None and node_name is not None:
|
||||
raise exceptions.InternalError("Cannot specify both node and node_name")
|
||||
|
||||
if node is not None:
|
||||
stmt = select(PubsubSub).where(PubsubSub.node_id == node.id)
|
||||
else:
|
||||
stmt = (
|
||||
select(PubsubSub)
|
||||
.join(PubsubNode, PubsubSub.node_id == PubsubNode.id)
|
||||
.options(joinedload(PubsubSub.node))
|
||||
.where(
|
||||
PubsubNode.profile_id == self.profiles[client.profile],
|
||||
)
|
||||
)
|
||||
|
||||
if service is not None:
|
||||
stmt = stmt.where(PubsubNode.service == service)
|
||||
if node_name is not None:
|
||||
stmt = stmt.where(PubsubNode.name == node_name)
|
||||
|
||||
if entity is not None:
|
||||
stmt = stmt.where(PubsubSub.subscriber == entity)
|
||||
|
||||
async with self.session() as session:
|
||||
result = await session.execute(stmt)
|
||||
|
||||
subscriptions = result.scalars().all()
|
||||
|
||||
if node is not None:
|
||||
for sub in subscriptions:
|
||||
sub.node = node
|
||||
|
||||
return subscriptions
|
||||
|
||||
@aio
|
||||
async def set_pubsub_subscriptions(
|
||||
self,
|
||||
node: PubsubNode,
|
||||
subscriptions: dict[jid.JID, SubscriptionState | None],
|
||||
) -> None:
|
||||
"""Set subscriptions for a node.
|
||||
|
||||
Uses upsert to insert new subscriptions or update existing ones.
|
||||
Subscriptions with value None are removed.
|
||||
|
||||
@param node: The node whose subscriptions to set.
|
||||
@param subscriptions: Mapping of subscriber JIDs to their new subscription states.
|
||||
If state is None, the subscription is removed.
|
||||
"""
|
||||
async with self.session() as session:
|
||||
async with session.begin():
|
||||
to_set = {
|
||||
subscriber: state
|
||||
for subscriber, state in subscriptions.items()
|
||||
if state is not None
|
||||
}
|
||||
to_delete = [
|
||||
subscriber
|
||||
for subscriber, state in subscriptions.items()
|
||||
if state is None
|
||||
]
|
||||
|
||||
# Upsert subscriptions
|
||||
if to_set:
|
||||
for subscriber, state in to_set.items():
|
||||
stmt = (
|
||||
insert(PubsubSub)
|
||||
.values(
|
||||
node_id=node.id,
|
||||
subscriber=subscriber,
|
||||
state=state,
|
||||
)
|
||||
.on_conflict_do_update(
|
||||
index_elements=(
|
||||
PubsubSub.node_id,
|
||||
PubsubSub.subscriber,
|
||||
),
|
||||
set_={PubsubSub.state: state},
|
||||
)
|
||||
)
|
||||
await session.execute(stmt)
|
||||
|
||||
# Delete subscriptions marked as None
|
||||
if to_delete:
|
||||
delete_stmt = delete(PubsubSub).where(
|
||||
PubsubSub.node_id == node.id,
|
||||
PubsubSub.subscriber.in_(to_delete),
|
||||
)
|
||||
await session.execute(delete_stmt)
|
||||
|
||||
await session.commit()
|
||||
|
||||
def _get_sqlite_path(self, path: List[Union[str, int]]) -> str:
|
||||
"""generate path suitable to query JSON element with SQLite"""
|
||||
return f"${''.join(f'[{p}]' if isinstance(p, int) else f'.{p}' for p in path)}"
|
||||
|
|
|
|||
Loading…
Reference in a new issue