improvement(search): generalize connector partition progress

Author: waleedlatif1

.github/workflows/test-build.yml

@@ -238,7 +238,7 @@ jobs:
           lib/knowledge/__integration__/connector-lifecycle-locks.integration.ts
           lib/knowledge/__integration__/connector-deferral.integration.ts
           lib/knowledge/__integration__/stored-document-recovery.integration.ts
-          lib/knowledge/__integration__/google-company-work.integration.ts
+          lib/knowledge/__integration__/connector-partition-work.integration.ts
           lib/knowledge/__integration__/listing-continuation.integration.ts
           lib/knowledge/__integration__/kb-block-search.integration.ts
           lib/core/outbox/service.integration.ts

apps/sim/connectors/google-drive/company-crawl.ts

@@ -6,7 +6,10 @@ import {
   GOOGLE_WORKSPACE_DRIVES_PAGE_SIZE,
   listGoogleWorkspaceDrives,
 } from '@/connectors/google-drive/workspace-drives'
-import type { GoogleCompanyCursorAdapter } from '@/connectors/google-workspace/company-work'
+import {
+  type GoogleCompanyCursorAdapter,
+  googleCompanyUserContextSchema,
+} from '@/connectors/google-workspace/company-work'
 import {
   GOOGLE_WORKSPACE_USERS_PAGE_SIZE,
   getGoogleWorkspaceUser,
@@ -19,15 +22,7 @@ import { parseOptionalUnlimitedSafeInteger } from '@/connectors/utils'
 const CURSOR_PREFIX = 'gdrive-company:v1:'
 const MAX_CURSOR_BYTES = 384 * 1024
 const cursorSchema = z.object({
-  users: z
-    .array(
-      z.object({
-        id: z.string().min(1).max(256),
-        email: z.string().email().max(254),
-        customerId: z.string().min(1).max(256),
-      })
-    )
-    .max(GOOGLE_WORKSPACE_USERS_PAGE_SIZE),
+  users: z.array(googleCompanyUserContextSchema).max(GOOGLE_WORKSPACE_USERS_PAGE_SIZE),
   nextUsersPageToken: z.string().min(1).max(8192).optional(),
   scope: z.discriminatedUnion('kind', [
     z.object({

apps/sim/connectors/google-workspace/company-crawl.ts

@@ -3,7 +3,10 @@ import { normalizeEmail } from '@sim/utils/string'
 import { z } from 'zod'
 import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
 import { GoogleApiError } from '@/connectors/google-workspace/api-errors'
-import type { GoogleCompanyCursorAdapter } from '@/connectors/google-workspace/company-work'
+import {
+  type GoogleCompanyCursorAdapter,
+  googleCompanyUserContextSchema,
+} from '@/connectors/google-workspace/company-work'
 import {
   GOOGLE_WORKSPACE_USERS_PAGE_SIZE,
   type GoogleWorkspaceUser,
@@ -29,15 +32,7 @@ const MAX_PROVIDER_CURSOR_BYTES = 256 * 1024
 const MAX_PAGE_DOCUMENTS = 2500
 const cursorSchema = z.object({
   provider: z.enum(['gmail', 'google_calendar']),
-  users: z
-    .array(
-      z.object({
-        id: z.string().min(1).max(256),
-        email: z.string().email().max(254),
-        customerId: z.string().min(1).max(256),
-      })
-    )
-    .max(GOOGLE_WORKSPACE_USERS_PAGE_SIZE),
+  users: z.array(googleCompanyUserContextSchema).max(GOOGLE_WORKSPACE_USERS_PAGE_SIZE),
   nextUsersPageToken: z.string().min(1).max(8192).optional(),
   providerCursor: z.string().min(1).max(MAX_PROVIDER_CURSOR_BYTES).optional(),
   listingFailures: listingFailuresSchema.optional(),

apps/sim/connectors/google-workspace/company-work.ts

@@ -1,3 +1,4 @@
+import { z } from 'zod'
 import type { GoogleWorkspaceUser } from '@/connectors/google-workspace/users'
 
 /** Directory identity and provider continuation are persisted separately from delegated credentials. */
@@ -11,3 +12,10 @@ export interface GoogleCompanyCursorAdapter {
   /** Retains the exact active page when adopting the former single-user-at-a-time cursor. */
   resume: (cursor: string) => GoogleCompanyUserWork[]
 }
+
+/** Validates Google Directory identity before persisted partition context is used for delegation. */
+export const googleCompanyUserContextSchema = z.object({
+  id: z.string().min(1).max(256),
+  email: z.string().email().max(254),
+  customerId: z.string().min(1).max(256),
+})