add useCount to DataSync for efficient count queries

IHP/DataSync/ControllerImpl.hs

@@ -187,6 +187,49 @@ buildMessageHandler ensureRLSEnabled installTableChangeTriggers sendJSON handleC
 
                     MVar.takeMVar close
 
+            handleMessage CreateCountSubscription { query, requestId } = do
+                ensureBelowSubscriptionsLimit
+
+                tableNameRLS <- ensureRLSEnabled query.table
+
+                subscriptionId <- UUID.nextRandom
+
+                -- Allocate the close handle as early as possible
+                -- to make DeleteDataSubscription calls succeed even when the CountSubscription is
+                -- not fully set up yet
+                close <- MVar.newEmptyMVar
+                atomicModifyIORef'' ?state (\state -> state |> modify #subscriptions (HashMap.insert subscriptionId close))
+
+                let (theQuery, theParams) = compileQueryWithRenamer (renamer query.table) query
+
+                let countQuery = "SELECT COUNT(*) FROM (" <> theQuery <> ") AS _inner"
+
+                let
+                    unpackResult :: [(Only Int)] -> Int
+                    unpackResult [(Only value)] = value
+                    unpackResult otherwise = error "DataSync.unpackResult: Expected INT, but got something else"
+
+                count <- unpackResult <$> sqlQueryWithRLS countQuery theParams
+                countRef <- newIORef count
+
+                installTableChangeTriggers tableNameRLS
+
+                let
+                    callback :: ChangeNotifications.ChangeNotification -> IO ()
+                    callback _ = do
+                        newCount <- unpackResult <$> sqlQueryWithRLS countQuery theParams
+                        lastCount <- readIORef countRef
+
+                        when (newCount /= count) (sendJSON DidChangeCount { subscriptionId, count = newCount })
+
+                let subscribe = PGListener.subscribeJSON (ChangeNotifications.channelName tableNameRLS) callback pgListener
+                let unsubscribe subscription = PGListener.unsubscribe subscription pgListener
+
+                Exception.bracket subscribe unsubscribe \channelSubscription -> do
+                    sendJSON DidCreateCountSubscription { subscriptionId, requestId, count }
+
+                    MVar.takeMVar close
+
             handleMessage DeleteDataSubscription { requestId, subscriptionId } = do
                 DataSyncReady { subscriptions } <- getState
                 case HashMap.lookup subscriptionId subscriptions of

IHP/DataSync/Types.hs

@@ -10,6 +10,7 @@ import Control.Concurrent.MVar as MVar
 data DataSyncMessage
     = DataSyncQuery { query :: !DynamicSQLQuery, requestId :: !Int, transactionId :: !(Maybe UUID) }
     | CreateDataSubscription { query :: !DynamicSQLQuery, requestId :: !Int }
+    | CreateCountSubscription { query :: !DynamicSQLQuery, requestId :: !Int }
     | DeleteDataSubscription { subscriptionId :: !UUID, requestId :: !Int }
     | CreateRecordMessage { table :: !Text, record :: !(HashMap Text Value), requestId :: !Int, transactionId :: !(Maybe UUID) }
     | CreateRecordsMessage { table :: !Text, records :: ![HashMap Text Value], requestId :: !Int, transactionId :: !(Maybe UUID) }
@@ -31,10 +32,12 @@ data DataSyncResponse
     | DataSyncError { requestId :: !Int, errorMessage :: !Text }
     | FailedToDecodeMessageError { errorMessage :: !Text }
     | DidCreateDataSubscription { requestId :: !Int, subscriptionId :: !UUID, result :: ![[Field]] }
+    | DidCreateCountSubscription { requestId :: !Int, subscriptionId :: !UUID, count :: !Int }
     | DidDeleteDataSubscription { requestId :: !Int, subscriptionId :: !UUID }
     | DidInsert { subscriptionId :: !UUID, record :: ![Field] }
     | DidUpdate { subscriptionId :: !UUID, id :: UUID, changeSet :: !Value }
     | DidDelete { subscriptionId :: !UUID, id :: !UUID }
+    | DidChangeCount { subscriptionId :: !UUID, count :: !Int }
     | DidCreateRecord { requestId :: !Int, record :: ![Field] } -- ^ Response to 'CreateRecordMessage'
     | DidCreateRecords { requestId :: !Int, records :: ![[Field]] } -- ^ Response to 'CreateRecordsMessage'
     | DidUpdateRecord { requestId :: !Int, record :: ![Field] } -- ^ Response to 'UpdateRecordMessage'

ihp-datasync-typescript/IHP/DataSync/TypeScript/Compiler.hs

@@ -259,6 +259,8 @@ declare module 'ihp-datasync/react' {
      */
     function useQuery<table extends TableName, result>(queryBuilder: QueryBuilder<table, result>, options?: DataSubscriptionOptions): Array<result> | null;
 
+    function useCount<table extends TableName>(queryBuilder: QueryBuilder<table, any>): number | null;
+
     /**
      * A version of `useQuery` when you only want to fetch a single record.
      * 

ihp-datasync-typescript/Test/Spec.hs

@@ -349,6 +349,8 @@ tests = do
                          */
                         function useQuery<table extends TableName, result>(queryBuilder: QueryBuilder<table, result>, options?: DataSubscriptionOptions): Array<result> | null;
 
+                        function useCount<table extends TableName>(queryBuilder: QueryBuilder<table, any>): number | null;
+
                         /**
                          * A version of `useQuery` when you only want to fetch a single record.
                          * 

lib/IHP/DataSync/react.js

@@ -1,4 +1,4 @@
-import React, { useState, useEffect, useContext, useSyncExternalStore } from 'react';
+import React, { useState, useEffect, useContext, useSyncExternalStore, useRef, useMemo } from 'react';
 import { DataSubscription, DataSyncController } from './ihp-datasync.js';
 
 // Most IHP apps never use this context because they use session cookies for auth.
@@ -99,4 +99,43 @@ export class DataSubscriptionStore {
             return subscription;
         }
     }
+}
+
+export function useCount(queryBuilder) {
+    const count = useRef(null);
+    const getSnapshot = useMemo(() => () => count.current, []);
+    const subscribe = useMemo(() => (onStoreChange) => {
+        const controller = DataSyncController.getInstance();
+        var isActive = true;
+        var subscriptionId = null;
+        const onMessage = (message) => {
+            if (message.tag === 'DidChangeCount' && message.subscriptionId === subscriptionId) {
+                count.current = message.count;
+                onStoreChange();
+            }
+        };
+        controller.sendMessage({ tag: 'CreateCountSubscription', query: queryBuilder.query  })
+            .then((response) => {
+                if (isActive) {
+                    subscriptionId = response.subscriptionId;
+                    count.current = response.count;
+                    onStoreChange();
+
+                    controller.addEventListener('message', onMessage);
+                } else {
+                    controller.sendMessage({ tag: 'DeleteDataSubscription', subscriptionId: response.subscriptionId });
+                }
+            })
+
+        return () => {
+            isActive = false;
+
+            if (subscriptionId) {
+                controller.sendMessage({ tag: 'DeleteDataSubscription', subscriptionId });
+            }
+            controller.removeEventListener('message', onMessage);
+        }
+    }, [JSON.stringify(queryBuilder.query)]);
+
+    return useSyncExternalStore(subscribe, getSnapshot);
 }