Skip to content

Commit 589daf1

Browse files
EhespEhesp
committed
feat(firestore): rework firestore queries
1 parent 3e34301 commit 589daf1

9 files changed

+723
-472
lines changed

‎packages/firestore/src/index.ts

+78-2
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,20 @@
1515
*
1616
*/
1717

18-
import { SnapshotListenOptions } from "firebase/firestore";
18+
import {
19+
DocumentData,
20+
DocumentReference,
21+
DocumentSnapshot,
22+
getDoc,
23+
getDocFromCache,
24+
getDocFromServer,
25+
getDocs,
26+
getDocsFromCache,
27+
getDocsFromServer,
28+
Query,
29+
QuerySnapshot,
30+
SnapshotListenOptions,
31+
} from "firebase/firestore";
1932

2033
export type GetSnapshotSource = "server" | "cache";
2134

@@ -35,6 +48,69 @@ export type WithIdField<D, F = void> = F extends string
3548
? D & { [key in F]: string }
3649
: D;
3750

51+
export async function getSnapshot<T>(
52+
ref: DocumentReference<T>,
53+
source?: GetSnapshotSource
54+
): Promise<DocumentSnapshot<T>> {
55+
let snapshot: DocumentSnapshot<T>;
56+
57+
if (source === "cache") {
58+
snapshot = await getDocFromCache(ref);
59+
} else if (source === "server") {
60+
snapshot = await getDocFromServer(ref);
61+
} else {
62+
snapshot = await getDoc(ref);
63+
}
64+
65+
return snapshot;
66+
}
67+
68+
export type NamedQueryPromise<T> = () => Promise<Query<T> | null>;
69+
70+
export type NamedQuery<T = DocumentData> = Query<T> | NamedQueryPromise<T>;
71+
72+
export type QueryType<T> = Query<T> | NamedQuery<T>;
73+
74+
export async function getQuerySnapshot<T>(
75+
query: Query<T>,
76+
source?: GetSnapshotSource
77+
): Promise<QuerySnapshot<T>> {
78+
let snapshot: QuerySnapshot<T>;
79+
80+
if (source === "cache") {
81+
snapshot = await getDocsFromCache(query);
82+
} else if (source === "server") {
83+
snapshot = await getDocsFromServer(query);
84+
} else {
85+
snapshot = await getDocs(query);
86+
}
87+
88+
return snapshot;
89+
}
90+
91+
function isNamedQuery<T>(query: QueryType<T>): query is NamedQuery<T> {
92+
return typeof query === "function";
93+
}
94+
95+
export async function resolveQuery<T>(query: QueryType<T>): Promise<Query<T>> {
96+
if (isNamedQuery(query)) {
97+
if (typeof query === "function") {
98+
// Firebase throws an error if the query doesn't exist.
99+
const resolved = await query();
100+
return resolved!;
101+
}
102+
103+
return query;
104+
}
105+
106+
return query;
107+
}
108+
38109
export * from "./useFirestoreDocument";
110+
export * from "./useFirestoreDocumentData";
111+
export * from "./useFirestoreInfiniteQuery";
112+
export * from "./useFirestoreInfiniteQueryData";
39113
export * from "./useFirestoreQuery";
40-
export * from "./useFirestoreMutation";
114+
export * from "./useFirestoreQueryData";
115+
export * from "./mutations";
116+
export * from "./namedQuery";

‎packages/firestore/src/useFirestoreMutation.ts renamed to ‎packages/firestore/src/mutations.ts

File renamed without changes.

‎packages/firestore/src/namedQuery.ts

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import {
2+
Firestore,
3+
Query,
4+
namedQuery as firestoreNamedQuery,
5+
} from "firebase/firestore";
6+
import { NamedQuery } from "./index";
7+
8+
const namedQueryCache: { [key: string]: Query } = {};
9+
10+
export function namedQuery<T>(
11+
firestore: Firestore,
12+
name: string
13+
): NamedQuery<T> {
14+
const key = `${firestore.app.name}:${name}`;
15+
16+
if (namedQueryCache[key]) {
17+
return namedQueryCache[key] as Query<T>;
18+
}
19+
20+
return () =>
21+
firestoreNamedQuery(firestore, name).then((query) => {
22+
if (query) {
23+
namedQueryCache[key] = query;
24+
return query as Query<T>;
25+
}
26+
27+
return null;
28+
});
29+
}

‎packages/firestore/src/useFirestoreDocument.ts

+52-154
Original file line numberDiff line numberDiff line change
@@ -22,41 +22,21 @@ import {
2222
QueryKey,
2323
UseQueryOptions,
2424
UseQueryResult,
25+
hashQueryKey,
2526
} from "react-query";
2627
import {
2728
DocumentData,
2829
DocumentReference,
2930
DocumentSnapshot,
30-
getDoc,
31-
getDocFromCache,
32-
getDocFromServer,
3331
onSnapshot,
34-
SnapshotOptions,
3532
Unsubscribe,
3633
FirestoreError,
3734
} from "firebase/firestore";
38-
import {
39-
GetSnapshotSource,
40-
UseFirestoreHookOptions,
41-
WithIdField,
42-
} from "./index";
43-
44-
async function getSnapshot<T>(
45-
ref: DocumentReference<T>,
46-
source?: GetSnapshotSource
47-
): Promise<DocumentSnapshot<T>> {
48-
let snapshot: DocumentSnapshot<T>;
49-
50-
if (source === "cache") {
51-
snapshot = await getDocFromCache(ref);
52-
} else if (source === "server") {
53-
snapshot = await getDocFromServer(ref);
54-
} else {
55-
snapshot = await getDoc(ref);
56-
}
35+
import { getSnapshot, UseFirestoreHookOptions } from "./index";
36+
import { Completer } from "../../utils/src";
5737

58-
return snapshot;
59-
}
38+
const counts: { [key: string]: number } = {};
39+
const subscriptions: { [key: string]: Unsubscribe } = {};
6040

6141
export function useFirestoreDocument<T = DocumentData, R = DocumentSnapshot<T>>(
6242
key: QueryKey,
@@ -68,152 +48,70 @@ export function useFirestoreDocument<T = DocumentData, R = DocumentSnapshot<T>>(
6848
>
6949
): UseQueryResult<R, FirestoreError> {
7050
const client = useQueryClient();
71-
const unsubscribe = useRef<Unsubscribe>();
51+
const completer = useRef<Completer<DocumentSnapshot<T>>>(new Completer());
7252

73-
useEffect(() => {
74-
return () => unsubscribe.current?.();
75-
}, []);
53+
const hashFn = useQueryOptions?.queryKeyHashFn || hashQueryKey;
54+
const hash = hashFn(key);
7655

77-
return useQuery<DocumentSnapshot<T>, FirestoreError, R>({
78-
...useQueryOptions,
79-
queryKey: useQueryOptions?.queryKey ?? key,
80-
staleTime:
81-
useQueryOptions?.staleTime ?? options?.subscribe ? Infinity : undefined,
82-
async queryFn() {
83-
unsubscribe.current?.();
56+
const isSubscription = !!options?.subscribe;
8457

85-
if (!options?.subscribe) {
86-
return getSnapshot(ref, options?.source);
87-
}
58+
useEffect(() => {
59+
if (!isSubscription) {
60+
getSnapshot(ref, options?.source)
61+
.then((snapshot) => {
62+
completer.current!.complete(snapshot);
63+
})
64+
.catch((error) => {
65+
completer.current!.reject(error);
66+
});
67+
}
68+
}, [isSubscription, hash, completer]);
8869

89-
let resolved = false;
70+
useEffect(() => {
71+
if (isSubscription) {
72+
counts[hash] ??= 0;
73+
counts[hash]++;
9074

91-
return new Promise<DocumentSnapshot<T>>((resolve, reject) => {
92-
unsubscribe.current = onSnapshot(
75+
// If there is only one instance of this query key, subscribe
76+
if (counts[hash] === 1) {
77+
subscriptions[hash] = onSnapshot(
9378
ref,
9479
{
9580
includeMetadataChanges: options?.includeMetadataChanges,
9681
},
9782
(snapshot) => {
98-
if (!resolved) {
99-
resolved = true;
100-
return resolve(snapshot);
101-
} else {
102-
client.setQueryData<DocumentSnapshot<T>>(key, snapshot);
83+
// Set the data each time state changes.
84+
client.setQueryData<DocumentSnapshot<T>>(key, snapshot);
85+
86+
// Resolve the completer with the current data.
87+
if (!completer.current!.completed) {
88+
completer.current!.complete(snapshot);
10389
}
10490
},
105-
reject
91+
(error) => completer.current!.reject(error)
10692
);
107-
});
108-
},
109-
});
110-
}
111-
112-
export function useFirestoreDocumentData<
113-
T = DocumentData,
114-
R = WithIdField<T> | undefined
115-
>(
116-
key: QueryKey,
117-
ref: DocumentReference<T>,
118-
options?: UseFirestoreHookOptions & SnapshotOptions,
119-
useQueryOptions?: Omit<
120-
UseQueryOptions<WithIdField<T> | undefined, FirestoreError, R>,
121-
"queryFn"
122-
>
123-
): UseQueryResult<R, FirestoreError>;
124-
125-
export function useFirestoreDocumentData<
126-
ID extends string,
127-
T = DocumentData,
128-
R = WithIdField<T, ID> | undefined
129-
>(
130-
key: QueryKey,
131-
ref: DocumentReference<T>,
132-
options?: UseFirestoreHookOptions & SnapshotOptions & { idField: ID },
133-
useQueryOptions?: Omit<
134-
UseQueryOptions<WithIdField<T, ID> | undefined, FirestoreError, R>,
135-
"queryFn"
136-
>
137-
): UseQueryResult<R | undefined, FirestoreError>;
93+
} else {
94+
// Since there is already an active subscription, resolve the completer
95+
// with the cached data.
96+
completer.current!.complete(
97+
client.getQueryData(key) as DocumentSnapshot<T>
98+
);
99+
}
138100

139-
export function useFirestoreDocumentData<
140-
ID extends string,
141-
T = DocumentData,
142-
R = WithIdField<T, ID> | undefined
143-
>(
144-
key: QueryKey,
145-
ref: DocumentReference<T>,
146-
options?: UseFirestoreHookOptions & SnapshotOptions & { idField?: ID },
147-
useQueryOptions?: Omit<
148-
UseQueryOptions<WithIdField<T, ID> | undefined, FirestoreError, R>,
149-
"queryFn"
150-
>
151-
): UseQueryResult<R, FirestoreError> {
152-
const client = useQueryClient();
153-
const unsubscribe = useRef<Unsubscribe>();
101+
return () => {
102+
counts[hash]--;
154103

155-
useEffect(() => {
156-
return () => unsubscribe.current?.();
157-
}, []);
104+
if (counts[hash] === 0) {
105+
subscriptions[hash]();
106+
delete subscriptions[hash];
107+
}
108+
};
109+
}
110+
}, [isSubscription, hash, completer]);
158111

159-
return useQuery<WithIdField<T, ID> | undefined, FirestoreError, R>({
112+
return useQuery<DocumentSnapshot<T>, FirestoreError, R>({
160113
...useQueryOptions,
161114
queryKey: useQueryOptions?.queryKey ?? key,
162-
staleTime:
163-
useQueryOptions?.staleTime ?? options?.subscribe ? Infinity : undefined,
164-
async queryFn(): Promise<WithIdField<T, ID> | undefined> {
165-
unsubscribe.current?.();
166-
167-
if (!options?.subscribe) {
168-
const snapshot = await getSnapshot(ref, options?.source);
169-
170-
let data = snapshot.data({
171-
serverTimestamps: options?.serverTimestamps,
172-
});
173-
174-
if (data && options?.idField) {
175-
data = {
176-
...data,
177-
[options.idField]: snapshot.id,
178-
};
179-
}
180-
181-
return data as WithIdField<T, ID> | undefined;
182-
}
183-
184-
let resolved = false;
185-
186-
return new Promise<WithIdField<T, ID> | undefined>((resolve, reject) => {
187-
unsubscribe.current = onSnapshot(
188-
ref,
189-
{
190-
includeMetadataChanges: options?.includeMetadataChanges,
191-
},
192-
(snapshot) => {
193-
let data = snapshot.data({
194-
serverTimestamps: options?.serverTimestamps,
195-
});
196-
197-
if (data && options?.idField) {
198-
data = {
199-
...data,
200-
[options.idField]: snapshot.id,
201-
};
202-
}
203-
204-
if (!resolved) {
205-
resolved = true;
206-
return resolve(data as WithIdField<T, ID> | undefined);
207-
} else {
208-
client.setQueryData<WithIdField<T, ID> | undefined>(
209-
key,
210-
data as WithIdField<T, ID> | undefined
211-
);
212-
}
213-
},
214-
reject
215-
);
216-
});
217-
},
115+
queryFn: () => completer.current!.promise,
218116
});
219117
}

0 commit comments

Comments
 (0)