fix: add ClearAll dataloader method

Signed-off-by: Matt Krick <[email protected]>

packages/server/billing/helpers/adjustUserCount.ts

@@ -82,7 +82,6 @@ const addUser = async (orgIds: string[], user: IUser, dataLoader: DataLoaderWork
     dataLoader.get('organizations').loadMany(orgIds),
     dataLoader.get('organizationUsersByUserId').load(userId)
   ])
-  dataLoader.get('organizationUsersByUserId').clear(userId)
   const organizations = rawOrganizations.filter(isValid)
   const docs = orgIds.map((orgId) => {
     const oldOrganizationUser = organizationUsers.find(
@@ -97,6 +96,7 @@ const addUser = async (orgIds: string[], user: IUser, dataLoader: DataLoaderWork
       tier: organization.tier
     })
   })
+  dataLoader.clearAll('organizationUsers')
   await getKysely()
     .insertInto('OrganizationUser')
     .values(docs)

packages/server/dataloader/RootDataLoader.ts

@@ -1,8 +1,6 @@
 import DataLoader from 'dataloader'
-import {DBType} from '../database/rethinkDriver'
 import RethinkForeignKeyLoaderMaker from './RethinkForeignKeyLoaderMaker'
 import RethinkPrimaryKeyLoaderMaker from './RethinkPrimaryKeyLoaderMaker'
-import UpdatableCacheDataLoader from './UpdatableCacheDataLoader'
 import * as atlassianLoaders from './atlassianLoaders'
 import * as azureDevOpsLoaders from './azureDevOpsLoaders'
 import * as customLoaderMakers from './customLoaderMakers'
@@ -41,79 +39,81 @@ const loaderMakers = {
   ...azureDevOpsLoaders
 } as const
 
-type LoaderMakers = typeof loaderMakers
-export type Loaders = keyof LoaderMakers
+export type Loaders = keyof typeof loaderMakers
 
-type PrimaryLoaderMakers = typeof rethinkPrimaryKeyLoaderMakers
-type PrimaryLoaders = keyof PrimaryLoaderMakers
-type Unprimary<T> = T extends RethinkPrimaryKeyLoaderMaker<infer U> ? DBType[U] : never
-type TypeFromPrimary<T extends PrimaryLoaders> = Unprimary<PrimaryLoaderMakers[T]>
+export type AllPrimaryLoaders =
+  | keyof typeof primaryKeyLoaderMakers
+  | keyof typeof rethinkPrimaryKeyLoaderMakers
+export type RegisterDependsOn = (primaryLoaders: AllPrimaryLoaders | AllPrimaryLoaders[]) => void
 
-type ForeignLoaderMakers = typeof rethinkForeignKeyLoaderMakers
-type ForeignLoaders = keyof ForeignLoaderMakers
-type Unforeign<T> = T extends RethinkForeignKeyLoaderMaker<infer U> ? U : never
-type TypeFromForeign<T extends ForeignLoaders> = TypeFromPrimary<Unforeign<ForeignLoaderMakers[T]>>
-
-/**
- * When adding a new loaders file like {@link atlassianLoaders} or {@link githubLoaders}
- * this type has to include a typeof of newly added loaders
- */
-type CustomLoaderMakers = typeof customLoaderMakers &
-  typeof atlassianLoaders &
-  typeof jiraServerLoaders &
-  typeof pollLoaders &
-  typeof integrationAuthLoaders &
-  typeof primaryKeyLoaderMakers &
-  typeof foreignKeyLoaderMakers &
-  typeof azureDevOpsLoaders &
-  typeof gitlabLoaders &
-  typeof gcalLoaders
-type CustomLoaders = keyof CustomLoaderMakers
-type Uncustom<T> = T extends (parent: RootDataLoader) => infer U ? U : never
-type TypeFromCustom<T extends CustomLoaders> = Uncustom<CustomLoaderMakers[T]>
-
-// Use this if you don't need the dataloader to be shareable
-export interface DataLoaderInstance {
-  get<LoaderName extends Loaders>(loaderName: LoaderName): TypedDataLoader<LoaderName>
+// The RethinkDB logic is a leaky abstraction! It will be gone soon & this will be generic enough to put in its own package
+interface GenericDataLoader<TLoaders, TPrimaryLoaderNames> {
+  clearAll(pkLoaderName: TPrimaryLoaderNames): void
+  get<LoaderName extends keyof TLoaders, Loader extends TLoaders[LoaderName]>(
+    loaderName: LoaderName
+  ): Loader extends (...args: any[]) => any
+    ? ReturnType<Loader>
+    : // can delete below this line after RethinkDB is gone
+      Loader extends RethinkPrimaryKeyLoaderMaker<infer U>
+      ? ReturnType<typeof rethinkPrimaryKeyLoader<U>>
+      : Loader extends RethinkForeignKeyLoaderMaker<infer U>
+        ? ReturnType<
+            typeof rethinkForeignKeyLoader<
+              (typeof rethinkPrimaryKeyLoaderMakers)[U] extends RethinkPrimaryKeyLoaderMaker<
+                infer V
+              >
+                ? V
+                : never
+            >
+          >
+        : never
 }
 
-export type TypedDataLoader<LoaderName> = LoaderName extends CustomLoaders
-  ? TypeFromCustom<LoaderName>
-  : UpdatableCacheDataLoader<
-      string,
-      LoaderName extends ForeignLoaders
-        ? TypeFromForeign<LoaderName>[]
-        : LoaderName extends PrimaryLoaders
-          ? TypeFromPrimary<LoaderName>
-          : never
-    >
+export type DataLoaderInstance = GenericDataLoader<typeof loaderMakers, AllPrimaryLoaders>
 
 /**
  * This is the main dataloader
  */
-export default class RootDataLoader implements DataLoaderInstance {
-  dataLoaderOptions: DataLoader.Options<any, any>
+export default class RootDataLoader<
+  O extends DataLoader.Options<any, any> = DataLoader.Options<any, any>
+> {
+  dataLoaderOptions: O
   // casted to any because access to the loaders will results in a creation if needed
   loaders: LoaderDict = {} as any
-  constructor(dataLoaderOptions: DataLoader.Options<any, any> = {}) {
+  dependentLoaders: Record<string, string[]> = {}
+  constructor(dataLoaderOptions: O = {} as O) {
     this.dataLoaderOptions = dataLoaderOptions
   }
 
-  get<LoaderName extends Loaders>(loaderName: LoaderName): TypedDataLoader<LoaderName> {
+  clearAll: DataLoaderInstance['clearAll'] = (pkLoaderName) => {
+    const dependencies = [pkLoaderName, ...(this.dependentLoaders[pkLoaderName] ?? [])]
+    dependencies.forEach((loaderName) => {
+      this.loaders[loaderName]?.clearAll()
+    })
+  }
+  get: DataLoaderInstance['get'] = (loaderName) => {
     let loader = this.loaders[loaderName]
-    if (loader) return loader as TypedDataLoader<LoaderName>
-    const loaderMaker = loaderMakers[loaderName]
+    if (loader) return loader
+    const loaderMaker = loaderMakers[loaderName as keyof typeof loaderMakers]
     if (loaderMaker instanceof RethinkPrimaryKeyLoaderMaker) {
       const {table} = loaderMaker
       loader = rethinkPrimaryKeyLoader(this.dataLoaderOptions, table)
     } else if (loaderMaker instanceof RethinkForeignKeyLoaderMaker) {
       const {fetch, field, pk} = loaderMaker
-      const basePkLoader = this.get(pk as PrimaryLoaders)
+      const basePkLoader = this.get(pk) as any
       loader = rethinkForeignKeyLoader(basePkLoader, this.dataLoaderOptions, field, fetch)
     } else {
-      loader = (loaderMaker as any)(this)
+      const dependsOn: RegisterDependsOn = (inPrimaryLoaders) => {
+        const primaryLoaders = Array.isArray(inPrimaryLoaders)
+          ? inPrimaryLoaders
+          : [inPrimaryLoaders]
+        primaryLoaders.forEach((primaryLoader) => {
+          ;(this.dependentLoaders[primaryLoader] ??= []).push(loaderName)
+        })
+      }
+      loader = (loaderMaker as any)(this, dependsOn)
     }
     this.loaders[loaderName] = loader!
-    return loader as TypedDataLoader<LoaderName>
+    return loader as any
   }
 }