batch assoc functions

Author: evgeny-simplex

src/Simplex/Messaging/Server/QueueStore/Postgres.hs

@@ -505,7 +505,28 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
         withLog "setQueueService" st $ \sl -> logQueueService sl rId party serviceId
 
   setQueueServices _ _ _ [] = pure $ Right M.empty
-  setQueueServices _ _ _ _ = pure $ Right M.empty -- TODO batch implementation
+  setQueueServices st party serviceId qs = E.uninterruptibleMask_ $ runExceptT $ do
+    updated <- S.fromList <$> withDB' "setQueueServices" st (\db ->
+      map fromOnly <$> DB.query db updateQuery (serviceId, In (map recipientId qs)))
+    forM_ qs $ \sq -> when (S.member (recipientId sq) updated) $ do
+      ExceptT $ readQueueRecIO (queueRec sq) >>= \case
+        Left e -> pure $ Left e
+        Right q -> runExceptT $ do
+          let q' = updateRec q
+          atomically $ writeTVar (queueRec sq) $ Just q'
+          withLog "setQueueServices" st $ \sl -> logQueueService sl (recipientId sq) party serviceId
+    pure $ M.fromList [(recipientId sq, if S.member (recipientId sq) updated then Right () else Left AUTH) | sq <- qs]
+    where
+      updateQuery = case party of
+        SRecipientService ->
+          "UPDATE msg_queues SET rcv_service_id = ? WHERE recipient_id IN ? AND deleted_at IS NULL RETURNING recipient_id"
+        SNotifierService ->
+          "UPDATE msg_queues SET ntf_service_id = ? WHERE recipient_id IN ? AND notifier_id IS NOT NULL AND deleted_at IS NULL RETURNING recipient_id"
+      updateRec q = case party of
+        SRecipientService -> q {rcvServiceId = serviceId}
+        SNotifierService -> case notifier q of
+          Just nc -> q {notifier = Just nc {ntfServiceId = serviceId}}
+          Nothing -> q
 
   getQueueNtfServices :: PostgresQueueStore q -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getQueueNtfServices st ntfs = E.uninterruptibleMask_ $ runExceptT $ do

src/Simplex/Messaging/Server/QueueStore/STM.hs

@@ -337,7 +337,9 @@ instance StoreQueueClass q => QueueStoreClass q (STMQueueStore q) where
         mapM_ (removeServiceQueue st serviceSel qId) prevSrvId
         mapM_ (addServiceQueue st serviceSel qId) serviceId
 
-  setQueueServices _ _ _ _ = pure $ Right M.empty -- TODO loop implementation
+  setQueueServices st party serviceId qs = Right . M.fromList <$> mapM setOne qs
+    where
+      setOne sq = (recipientId sq,) <$> setQueueService st sq party serviceId
 
   getQueueNtfServices :: STMQueueStore q -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getQueueNtfServices st ntfs = do