handle a mix of requested info

gajus · May 13, 2024 · d20e205 · d20e205
1 parent 4b4f40c
commit d20e205
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 4 deletions.
diff --git a/packages/slonik-dataloaders/src/factories/createConnectionLoaderClass.test.ts b/packages/slonik-dataloaders/src/factories/createConnectionLoaderClass.test.ts
@@ -337,6 +337,17 @@ describe('createConnectionLoaderClass', () => {
   });
 
   it('gets the count without fetching edges', async () => {
+    const loader = new PersonConnectionLoader(pool, {});
+    const result = await loader.load({
+      info: getInfo(['count']),
+      where: ({ name }) => sql.fragment`${name} = 'ccc'`,
+    });
+
+    expect(result.count).toEqual(2n);
+    expect(result.edges.length).toEqual(0);
+  });
+
+  it('gets the count without fetching edges (batch)', async () => {
     const loader = new PersonConnectionLoader(pool, {});
     const results = await Promise.all([
       loader.load({
@@ -355,6 +366,24 @@ describe('createConnectionLoaderClass', () => {
     expect(results[1].edges.length).toEqual(0);
   });
 
+  it('gets a mix of count and edges', async () => {
+    const loader = new PersonConnectionLoader(pool, {});
+
+    const results = await Promise.all([
+      loader.load({
+        info: getInfo(['edges']),
+        where: ({ name }) => sql.fragment`${name} = 'eee'`,
+      }),
+      loader.load({
+        info: getInfo(['count']),
+        where: ({ name }) => sql.fragment`${name} = 'eee'`,
+      }),
+    ]);
+
+    expect(results[0].count).toEqual(0);
+    expect(results[1].count).toEqual(2n);
+  });
+
   it('gets the edges without fetching edges', async () => {
     const loader = new PersonConnectionLoader(pool, {});
     const results = await Promise.all([

diff --git a/packages/slonik-dataloaders/src/factories/createConnectionLoaderClass.ts b/packages/slonik-dataloaders/src/factories/createConnectionLoaderClass.ts
@@ -38,6 +38,7 @@ export const createConnectionLoaderClass = <T extends ZodTypeAny>(config: {
   query: QuerySqlToken<T>;
 }) => {
   const { columnNameTransformer = snakeCase, query } = config;
+
   const columnIdentifiers = getColumnIdentifiers<z.infer<T>>(
     TABLE_ALIAS,
     columnNameTransformer,
@@ -58,8 +59,8 @@ export const createConnectionLoaderClass = <T extends ZodTypeAny>(config: {
     ) {
       super(
         async (loaderKeys) => {
-          const edgesQueries: QuerySqlToken[] = [];
-          const countQueries: QuerySqlToken[] = [];
+          const edgesQueries: Array<QuerySqlToken | null> = [];
+          const countQueries: Array<QuerySqlToken | null> = [];
 
           for (const loaderKey of loaderKeys.values()) {
             const {
@@ -98,6 +99,8 @@ export const createConnectionLoaderClass = <T extends ZodTypeAny>(config: {
                   }
                 )`,
               );
+            } else {
+              countQueries.push(null);
             }
 
             if (
@@ -181,6 +184,8 @@ export const createConnectionLoaderClass = <T extends ZodTypeAny>(config: {
                   OFFSET ${offset || 0}
                 )`,
               );
+            } else {
+              edgesQueries.push(null);
             }
           }
 
@@ -203,12 +208,16 @@ export const createConnectionLoaderClass = <T extends ZodTypeAny>(config: {
           const [edgeResults, countResults] = await Promise.all([
             Promise.all(
               edgesQueries.map((edgesQuery) => {
-                return pool.any(sql.type(edgeSchema)`${edgesQuery}`);
+                return edgesQuery === null
+                  ? []
+                  : pool.any(sql.type(edgeSchema)`${edgesQuery}`);
               }),
             ),
             Promise.all(
               countQueries.map((countQuery) => {
-                return pool.oneFirst(sql.type(countSchema)`${countQuery}`);
+                return countQuery === null
+                  ? 0
+                  : pool.oneFirst(sql.type(countSchema)`${countQuery}`);
               }),
             ),
           ]);