chore(rethinkdb): Organization: Phase 1 (#9883)

Signed-off-by: Matt Krick <matt.krick@gmail.com>

packages/server/billing/helpers/adjustUserCount.ts

@@ -1,9 +1,10 @@
 import {InvoiceItemType} from 'parabol-client/types/constEnums'
 import getRethink from '../../database/rethinkDriver'
 import {RDatum} from '../../database/stricterR'
-import Organization from '../../database/types/Organization'
 import OrganizationUser from '../../database/types/OrganizationUser'
 import {DataLoaderWorker} from '../../graphql/graphql'
+import isValid from '../../graphql/isValid'
+import getKysely from '../../postgres/getKysely'
 import insertOrgUserAudit from '../../postgres/helpers/insertOrgUserAudit'
 import {OrganizationUserAuditEventTypeEnum} from '../../postgres/queries/generated/insertOrgUserAuditQuery'
 import {getUserById} from '../../postgres/queries/getUsersByIds'
@@ -22,7 +23,7 @@ const maybeUpdateOrganizationActiveDomain = async (
   dataLoader: DataLoaderWorker
 ) => {
   const r = await getRethink()
-  const organization = await r.table('Organization').get(orgId).run()
+  const organization = await dataLoader.get('organizations').load(orgId)
   const {isActiveDomainTouched, activeDomain} = organization
   // don't modify if the domain was set manually
   if (isActiveDomainTouched) return
@@ -38,14 +39,18 @@ const maybeUpdateOrganizationActiveDomain = async (
   // don't modify if we can't guess the domain or the domain we guess is the current domain
   const domain = await getActiveDomainForOrgId(orgId)
   if (!domain || domain === activeDomain) return
-
-  await r
-    .table('Organization')
-    .get(orgId)
-    .update({
-      activeDomain: domain
-    })
-    .run()
+  organization.activeDomain = domain
+  const pg = getKysely()
+  await Promise.all([
+    pg.updateTable('Organization').set({activeDomain: domain}).where('id', '=', orgId).execute(),
+    r
+      .table('Organization')
+      .get(orgId)
+      .update({
+        activeDomain: domain
+      })
+      .run()
+  ])
 }
 
 const changePause = (inactive: boolean) => async (_orgIds: string[], user: IUser) => {
@@ -76,18 +81,16 @@ const changePause = (inactive: boolean) => async (_orgIds: string[], user: IUser
 const addUser = async (orgIds: string[], user: IUser, dataLoader: DataLoaderWorker) => {
   const {id: userId} = user
   const r = await getRethink()
-  const {organizations, organizationUsers} = await r({
-    organizationUsers: r
+  const [rawOrganizations, organizationUsers] = await Promise.all([
+    dataLoader.get('organizations').loadMany(orgIds),
+    r
       .table('OrganizationUser')
       .getAll(userId, {index: 'userId'})
       .orderBy(r.desc('newUserUntil'))
-      .coerceTo('array') as unknown as OrganizationUser[],
-    organizations: r
-      .table('Organization')
-      .getAll(r.args(orgIds))
-      .coerceTo('array') as unknown as Organization[]
-  }).run()
-
+      .coerceTo('array')
+      .run()
+  ])
+  const organizations = rawOrganizations.filter(isValid)
   const docs = orgIds.map((orgId) => {
     const oldOrganizationUser = organizationUsers.find(
       (organizationUser) => organizationUser.orgId === orgId
@@ -153,7 +156,6 @@ export default async function adjustUserCount(
   type: InvoiceItemType,
   dataLoader: DataLoaderWorker
 ) {
-  const r = await getRethink()
   const orgIds = Array.isArray(orgInput) ? orgInput : [orgInput]
 
   const user = (await getUserById(userId))!
@@ -164,11 +166,8 @@ export default async function adjustUserCount(
   const auditEventType = auditEventTypeLookup[type]
   await insertOrgUserAudit(orgIds, userId, auditEventType)
 
-  const paidOrgs = await r
-    .table('Organization')
-    .getAll(r.args(orgIds), {index: 'id'})
-    .filter((org: RDatum) => org('stripeSubscriptionId').default(null).ne(null))
-    .run()
+  const organizations = await dataLoader.get('organizations').loadMany(orgIds)
+  const paidOrgs = organizations.filter(isValid).filter((org) => org.stripeSubscriptionId)
 
   handleEnterpriseOrgQuantityChanges(paidOrgs, dataLoader).catch()
   handleTeamOrgQuantityChanges(paidOrgs).catch(Logger.error)

packages/server/billing/helpers/generateInvoice.ts

@@ -8,7 +8,6 @@ import {InvoiceLineItemEnum} from '../../database/types/InvoiceLineItem'
 import InvoiceLineItemDetail from '../../database/types/InvoiceLineItemDetail'
 import InvoiceLineItemOtherAdjustments from '../../database/types/InvoiceLineItemOtherAdjustments'
 import NextPeriodCharges from '../../database/types/NextPeriodCharges'
-import Organization from '../../database/types/Organization'
 import QuantityChangeLineItem from '../../database/types/QuantityChangeLineItem'
 import generateUID from '../../generateUID'
 import {DataLoaderWorker} from '../../graphql/graphql'
@@ -354,16 +353,16 @@ export default async function generateInvoice(
     invoice.status === 'paid' && invoice.status_transitions.paid_at
       ? fromEpochSeconds(invoice.status_transitions.paid_at)
       : undefined
-
-  const {organization, billingLeaderIds} = await r({
-    organization: r.table('Organization').get(orgId) as unknown as Organization,
-    billingLeaderIds: r
+  const [organization, billingLeaderIds] = await Promise.all([
+    dataLoader.get('organizations').load(orgId),
+    r
       .table('OrganizationUser')
       .getAll(orgId, {index: 'orgId'})
       .filter({removedAt: null})
       .filter((row: RDatum) => r.expr(['BILLING_LEADER', 'ORG_ADMIN']).contains(row('role')))
-      .coerceTo('array')('userId') as unknown as string[]
-  }).run()
+      .coerceTo('array')('userId')
+      .run() as any as string[]
+  ])
 
   const billingLeaders = (await dataLoader.get('users').loadMany(billingLeaderIds)).filter(isValid)
   const billingLeaderEmails = billingLeaders.map((user) => user.email)

packages/server/billing/helpers/generateUpcomingInvoice.ts

@@ -1,14 +1,13 @@
-import getRethink from '../../database/rethinkDriver'
 import {DataLoaderWorker} from '../../graphql/graphql'
 import getUpcomingInvoiceId from '../../utils/getUpcomingInvoiceId'
 import {getStripeManager} from '../../utils/stripe'
 import fetchAllLines from './fetchAllLines'
 import generateInvoice from './generateInvoice'
 
 const generateUpcomingInvoice = async (orgId: string, dataLoader: DataLoaderWorker) => {
-  const r = await getRethink()
   const invoiceId = getUpcomingInvoiceId(orgId)
-  const {stripeId} = await r.table('Organization').get(orgId).pluck('stripeId').run()
+  const organization = await dataLoader.get('organizations').load(orgId)
+  const {stripeId} = organization
   const manager = getStripeManager()
   const [stripeLineItems, upcomingInvoice] = await Promise.all([
     fetchAllLines('upcoming', stripeId),

packages/server/billing/helpers/teamLimitsCheck.ts

@@ -80,7 +80,13 @@ export const maybeRemoveRestrictions = async (orgId: string, dataLoader: DataLoa
 
   if (!(await isLimitExceeded(orgId))) {
     const billingLeadersIds = await dataLoader.get('billingLeadersIdsByOrgId').load(orgId)
+    const pg = getKysely()
     await Promise.all([
+      pg
+        .updateTable('Organization')
+        .set({tierLimitExceededAt: null, scheduledLockAt: null, lockedAt: null})
+        .where('id', '=', orgId)
+        .execute(),
       r
         .table('Organization')
         .get(orgId)
@@ -129,16 +135,26 @@ export const checkTeamsLimit = async (orgId: string, dataLoader: DataLoaderWorke
 
   const now = new Date()
   const scheduledLockAt = new Date(now.getTime() + ms(`${Threshold.STARTER_TIER_LOCK_AFTER_DAYS}d`))
-
-  await r
-    .table('Organization')
-    .get(orgId)
-    .update({
-      tierLimitExceededAt: now,
-      scheduledLockAt,
-      updatedAt: now
-    })
-    .run()
+  const pg = getKysely()
+  await Promise.all([
+    pg
+      .updateTable('Organization')
+      .set({
+        tierLimitExceededAt: now,
+        scheduledLockAt
+      })
+      .where('id', '=', orgId)
+      .execute(),
+    r
+      .table('Organization')
+      .get(orgId)
+      .update({
+        tierLimitExceededAt: now,
+        scheduledLockAt,
+        updatedAt: now
+      })
+      .run()
+  ])
   dataLoader.get('organizations').clear(orgId)
 
   const billingLeaders = await getBillingLeadersByOrgId(orgId, dataLoader)

packages/server/billing/helpers/terminateSubscription.ts

@@ -1,31 +1,45 @@
 import getRethink from '../../database/rethinkDriver'
 import Organization from '../../database/types/Organization'
+import getKysely from '../../postgres/getKysely'
 import {Logger} from '../../utils/Logger'
+import sendToSentry from '../../utils/sendToSentry'
 import {getStripeManager} from '../../utils/stripe'
 
 const terminateSubscription = async (orgId: string) => {
   const r = await getRethink()
+  const pg = getKysely()
   const now = new Date()
   // flag teams as unpaid
-  const [rethinkResult] = await Promise.all([
-    r({
-      organization: r
-        .table('Organization')
-        .get(orgId)
-        .update(
-          {
-            // periodEnd should always be redundant, but useful for testing purposes
-            periodEnd: now,
-            stripeSubscriptionId: null
-          },
-          {returnChanges: true}
-        )('changes')(0)('old_val')
-        .default(null) as unknown as Organization
-    }).run()
+  const [pgOrganization, organization] = await Promise.all([
+    pg
+      .with('OldOrg', (qc) =>
+        qc.selectFrom('Organization').select('stripeSubscriptionId').where('id', '=', orgId)
+      )
+      .updateTable('Organization')
+      .set({periodEnd: now, stripeSubscriptionId: null})
+      .where('id', '=', orgId)
+      .returning((qc) =>
+        qc.selectFrom('OldOrg').select('stripeSubscriptionId').as('stripeSubscriptionId')
+      )
+      .executeTakeFirst(),
+    r
+      .table('Organization')
+      .get(orgId)
+      .update(
+        {
+          // periodEnd should always be redundant, but useful for testing purposes
+          periodEnd: now,
+          stripeSubscriptionId: null
+        },
+        {returnChanges: true}
+      )('changes')(0)('old_val')
+      .default(null)
+      .run() as unknown as Organization
   ])
-  const {organization} = rethinkResult
   const {stripeSubscriptionId} = organization
-
+  if (stripeSubscriptionId !== pgOrganization?.stripeSubscriptionId) {
+    sendToSentry(new Error(`stripeSubscriptionId mismatch for orgId ${orgId}`))
+  }
   if (stripeSubscriptionId) {
     const manager = getStripeManager()
     try {

packages/server/database/types/GoogleAnalyzedEntity.ts

@@ -1,5 +1,3 @@
-import {sql} from 'kysely'
-
 interface Input {
   lemma?: string
   name: string
@@ -17,8 +15,3 @@ export default class GoogleAnalyzedEntity {
     this.salience = salience
   }
 }
-
-export const toGoogleAnalyzedEntityPG = (entities: GoogleAnalyzedEntity[]) =>
-  sql<
-    string[]
-  >`(select coalesce(array_agg((name, salience, lemma)::"GoogleAnalyzedEntity"), '{}') from json_populate_recordset(null::"GoogleAnalyzedEntity", ${JSON.stringify(entities)}))`

packages/server/database/types/processTeamsLimitsJob.ts

@@ -3,6 +3,7 @@ import sendTeamsLimitEmail from '../../billing/helpers/sendTeamsLimitEmail'
 import {DataLoaderWorker} from '../../graphql/graphql'
 import isValid from '../../graphql/isValid'
 import publishNotification from '../../graphql/public/mutations/helpers/publishNotification'
+import getKysely from '../../postgres/getKysely'
 import NotificationTeamsLimitReminder from './NotificationTeamsLimitReminder'
 import ScheduledTeamLimitsJob from './ScheduledTeamLimitsJob'
 
@@ -27,7 +28,14 @@ const processTeamsLimitsJob = async (job: ScheduledTeamLimitsJob, dataLoader: Da
 
   if (type === 'LOCK_ORGANIZATION') {
     const now = new Date()
-    await r.table('Organization').get(orgId).update({lockedAt: now}).run()
+    await Promise.all([
+      getKysely()
+        .updateTable('Organization')
+        .set({lockedAt: now})
+        .where('id', '=', 'orgId')
+        .execute(),
+      r.table('Organization').get(orgId).update({lockedAt: now}).run()
+    ])
     organization.lockedAt = lockedAt
   } else if (type === 'WARN_ORGANIZATION') {
     const notificationsToInsert = billingLeadersIds.map((userId) => {

packages/server/graphql/mutations/addOrg.ts

@@ -62,7 +62,11 @@ export default {
     const teamId = generateUID()
     const {email} = viewer
     await createNewOrg(orgId, orgName, viewerId, email, dataLoader)
-    await createTeamAndLeader(viewer, {id: teamId, orgId, isOnboardTeam: false, ...newTeam})
+    await createTeamAndLeader(
+      viewer,
+      {id: teamId, orgId, isOnboardTeam: false, ...newTeam},
+      dataLoader
+    )
 
     const {tms} = authToken
     // MUTATIVE

packages/server/graphql/mutations/addTeam.ts

@@ -85,7 +85,7 @@ export default {
 
       // RESOLUTION
       const teamId = generateUID()
-      await createTeamAndLeader(viewer, {id: teamId, isOnboardTeam: false, ...newTeam})
+      await createTeamAndLeader(viewer, {id: teamId, isOnboardTeam: false, ...newTeam}, dataLoader)
 
       const {tms} = authToken
       // MUTATIVE

packages/server/graphql/mutations/createReflection.ts

@@ -6,10 +6,10 @@ import getGroupSmartTitle from 'parabol-client/utils/smartGroup/getGroupSmartTit
 import unlockAllStagesForPhase from 'parabol-client/utils/unlockAllStagesForPhase'
 import normalizeRawDraftJS from 'parabol-client/validation/normalizeRawDraftJS'
 import getRethink from '../../database/rethinkDriver'
-import {toGoogleAnalyzedEntityPG} from '../../database/types/GoogleAnalyzedEntity'
 import ReflectionGroup from '../../database/types/ReflectionGroup'
 import generateUID from '../../generateUID'
 import getKysely from '../../postgres/getKysely'
+import {toGoogleAnalyzedEntity} from '../../postgres/helpers/toGoogleAnalyzedEntity'
 import {analytics} from '../../utils/analytics/analytics'
 import {getUserId} from '../../utils/authorization'
 import publish from '../../utils/publish'
@@ -102,7 +102,7 @@ export default {
     await pg
       .with('Group', (qc) => qc.insertInto('RetroReflectionGroup').values(reflectionGroup))
       .insertInto('RetroReflection')
-      .values({...reflection, entities: toGoogleAnalyzedEntityPG(entities)})
+      .values({...reflection, entities: toGoogleAnalyzedEntity(entities)})
       .execute()
 
     const groupPhase = phases.find((phase) => phase.phaseType === 'group')!

packages/server/graphql/mutations/downgradeToStarter.ts

@@ -1,6 +1,5 @@
 import {GraphQLID, GraphQLList, GraphQLNonNull, GraphQLString} from 'graphql'
 import {SubscriptionChannel} from 'parabol-client/types/constEnums'
-import getRethink from '../../database/rethinkDriver'
 import {getUserId, isSuperUser, isUserBillingLeader} from '../../utils/authorization'
 import publish from '../../utils/publish'
 import standardError from '../../utils/standardError'
@@ -37,7 +36,6 @@ export default {
     }: {orgId: string; reasonsForLeaving?: TReasonToDowngradeEnum[]; otherTool?: string},
     {authToken, dataLoader, socketId: mutatorId}: GQLContext
   ) {
-    const r = await getRethink()
     const operationId = dataLoader.share()
     const subOptions = {mutatorId, operationId}
 
@@ -56,7 +54,8 @@ export default {
       return standardError(new Error('Other tool name is too long'), {userId: viewerId})
     }
 
-    const {stripeSubscriptionId, tier} = await r.table('Organization').get(orgId).run()
+    const {stripeSubscriptionId, tier} = await dataLoader.get('organizations').load(orgId)
+    dataLoader.get('organizations').clear(orgId)
 
     if (tier === 'starter') {
       return standardError(new Error('Already on free tier'), {userId: viewerId})
@@ -68,6 +67,7 @@ export default {
       orgId,
       stripeSubscriptionId!,
       viewer,
+      dataLoader,
       reasonsForLeaving,
       otherTool
     )

packages/server/graphql/mutations/helpers/bootstrapNewUser.ts

@@ -128,7 +128,7 @@ const bootstrapNewUser = async (
     const orgName = `${newUser.preferredName}’s Org`
     await createNewOrg(orgId, orgName, userId, email, dataLoader)
     await Promise.all([
-      createTeamAndLeader(newUser as IUser, validNewTeam),
+      createTeamAndLeader(newUser as IUser, validNewTeam, dataLoader),
       addSeedTasks(userId, teamId),
       r.table('SuggestedAction').insert(new SuggestedActionInviteYourTeam({userId, teamId})).run(),
       sendPromptToJoinOrg(newUser, dataLoader)