Skip to content

use useSyncExternalStore for useLiveQuery #225

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 6, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions packages/db/src/query/live-query-collection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,11 @@ export interface LiveQueryCollectionConfig<
* Start sync / the query immediately
*/
startSync?: boolean

/**
* GC time for the collection
*/
gcTime?: number
}

/**
Expand Down Expand Up @@ -322,6 +327,7 @@ export function liveQueryCollectionOptions<
config.getKey || ((item) => resultKeys.get(item) as string | number),
sync,
compare,
gcTime: config.gcTime || 5000, // 5 seconds by default for live queries
schema: config.schema,
onInsert: config.onInsert,
onUpdate: config.onUpdate,
Expand Down
137 changes: 100 additions & 37 deletions packages/react-db/src/useLiveQuery.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { useEffect, useMemo, useState } from "react"
import { useRef, useSyncExternalStore } from "react"
import { createLiveQueryCollection } from "@tanstack/db"
import type {
Collection,
Expand Down Expand Up @@ -55,58 +55,121 @@ export function useLiveQuery(
typeof configOrQueryOrCollection.startSyncImmediate === `function` &&
typeof configOrQueryOrCollection.id === `string`

const collection = useMemo(
() => {
if (isCollection) {
// It's already a collection, ensure sync is started for React hooks
configOrQueryOrCollection.startSyncImmediate()
return configOrQueryOrCollection
}
// Use refs to cache collection and track dependencies
const collectionRef = useRef<any>(null)
const depsRef = useRef<Array<unknown> | null>(null)
const configRef = useRef<any>(null)

// Check if we need to create/recreate the collection
const needsNewCollection =
!collectionRef.current ||
(isCollection && configRef.current !== configOrQueryOrCollection) ||
(!isCollection &&
(depsRef.current === null ||
depsRef.current.length !== deps.length ||
depsRef.current.some((dep, i) => dep !== deps[i])))

if (needsNewCollection) {
if (isCollection) {
// It's already a collection, ensure sync is started for React hooks
configOrQueryOrCollection.startSyncImmediate()
collectionRef.current = configOrQueryOrCollection
configRef.current = configOrQueryOrCollection
} else {
// Original logic for creating collections
// Ensure we always start sync for React hooks
if (typeof configOrQueryOrCollection === `function`) {
return createLiveQueryCollection({
collectionRef.current = createLiveQueryCollection({
query: configOrQueryOrCollection,
startSync: true,
gcTime: 0, // Live queries created by useLiveQuery are cleaned up immediately
})
} else {
return createLiveQueryCollection({
...configOrQueryOrCollection,
collectionRef.current = createLiveQueryCollection({
startSync: true,
gcTime: 0, // Live queries created by useLiveQuery are cleaned up immediately
...configOrQueryOrCollection,
})
}
},
isCollection ? [configOrQueryOrCollection] : [...deps]
)
depsRef.current = [...deps]
}
}

// Infer types from the actual collection
type CollectionType =
typeof collection extends Collection<infer T, any, any> ? T : never
type KeyType =
typeof collection extends Collection<any, infer K, any>
? K
: string | number
// Use refs to track version and memoized snapshot
const versionRef = useRef(0)
const snapshotRef = useRef<{
state: Map<any, any>
data: Array<any>
collection: Collection<any, any, any>
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we still infer the collection types?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, they are inferred in the function overloads, this is just internal and the type declaration was redundant.

_version: number
} | null>(null)

// Use a simple counter to force re-renders when collection changes
const [, forceUpdate] = useState(0)
// Reset refs when collection changes
if (needsNewCollection) {
versionRef.current = 0
snapshotRef.current = null
}

useEffect(() => {
// Subscribe to changes and force re-render
const unsubscribe = collection.subscribeChanges(() => {
forceUpdate((prev) => prev + 1)
})
// Create stable subscribe function using ref
const subscribeRef = useRef<
((onStoreChange: () => void) => () => void) | null
>(null)
if (!subscribeRef.current || needsNewCollection) {
subscribeRef.current = (onStoreChange: () => void) => {
const unsubscribe = collectionRef.current!.subscribeChanges(() => {
versionRef.current += 1
onStoreChange()
})
return () => {
unsubscribe()
}
}
}

// Create stable getSnapshot function using ref
const getSnapshotRef = useRef<
| (() => {
state: Map<any, any>
data: Array<any>
collection: Collection<any, any, any>
})
| null
>(null)
if (!getSnapshotRef.current || needsNewCollection) {
getSnapshotRef.current = () => {
const currentVersion = versionRef.current
const currentCollection = collectionRef.current!

// If we don't have a snapshot or the version changed, create a new one
if (
!snapshotRef.current ||
snapshotRef.current._version !== currentVersion
) {
snapshotRef.current = {
get state() {
return new Map(currentCollection.entries())
},
get data() {
return Array.from(currentCollection.values())
},
collection: currentCollection,
_version: currentVersion,
}
}

return unsubscribe
}, [collection])
return snapshotRef.current
}
}

// Use useSyncExternalStore to subscribe to collection changes
const snapshot = useSyncExternalStore(
subscribeRef.current,
getSnapshotRef.current
)

return {
get state(): Map<KeyType, CollectionType> {
return new Map(collection.entries())
},
get data(): Array<CollectionType> {
return Array.from(collection.values())
},
collection,
state: snapshot.state,
data: snapshot.data,
collection: snapshot.collection,
}
}