@@ -25,7 +25,7 @@ import { resolve } from 'path'
2525import { sleep } from '@sim/utils/helpers'
2626import { generateId } from '@sim/utils/id'
2727import { eq , sql } from 'drizzle-orm'
28- import { index , json , jsonb , pgTable , text , timestamp , uniqueIndex } from 'drizzle-orm/pg-core'
28+ import { index , json , jsonb , pgTable , text , timestamp } from 'drizzle-orm/pg-core'
2929import { drizzle } from 'drizzle-orm/postgres-js'
3030import postgres from 'postgres'
3131
@@ -196,7 +196,7 @@ const workspaceBYOKKeys = pgTable(
196196 updatedAt : timestamp ( 'updated_at' ) . notNull ( ) . defaultNow ( ) ,
197197 } ,
198198 ( table ) => ( {
199- workspaceProviderUnique : uniqueIndex ( 'workspace_byok_provider_unique ') . on (
199+ workspaceProviderIdx : index ( 'workspace_byok_workspace_provider_idx ') . on (
200200 table . workspaceId ,
201201 table . providerId
202202 ) ,
@@ -497,14 +497,12 @@ async function processWorkspace(
497497 stats . workspacesProcessed ++
498498
499499 const existingBYOKProviders = new Set < string > ( )
500- if ( DRY_RUN ) {
501- const existingRows = await db
502- . select ( { providerId : workspaceBYOKKeys . providerId } )
503- . from ( workspaceBYOKKeys )
504- . where ( eq ( workspaceBYOKKeys . workspaceId , workspaceId ) )
505- for ( const row of existingRows ) {
506- existingBYOKProviders . add ( row . providerId )
507- }
500+ const existingRows = await db
501+ . select ( { providerId : workspaceBYOKKeys . providerId } )
502+ . from ( workspaceBYOKKeys )
503+ . where ( eq ( workspaceBYOKKeys . workspaceId , workspaceId ) )
504+ for ( const row of existingRows ) {
505+ existingBYOKProviders . add ( row . providerId )
508506 }
509507
510508 let hasNewInserts = false
@@ -565,29 +563,25 @@ async function processWorkspace(
565563 continue
566564 }
567565
566+ if ( existingBYOKProviders . has ( providerId ) ) {
567+ console . log ( ` [SKIP] BYOK already exists for provider "${ providerId } "` )
568+ stats . skippedExisting ++
569+ continue
570+ }
571+
568572 try {
569573 const encrypted = await encryptSecret ( chosen . key )
570- const result = await db
571- . insert ( workspaceBYOKKeys )
572- . values ( {
573- id : generateId ( ) ,
574- workspaceId,
575- providerId,
576- encryptedApiKey : encrypted ,
577- createdBy : chosen . ref . userId ,
578- } )
579- . onConflictDoNothing ( {
580- target : [ workspaceBYOKKeys . workspaceId , workspaceBYOKKeys . providerId ] ,
581- } )
582- . returning ( { id : workspaceBYOKKeys . id } )
574+ await db . insert ( workspaceBYOKKeys ) . values ( {
575+ id : generateId ( ) ,
576+ workspaceId,
577+ providerId,
578+ encryptedApiKey : encrypted ,
579+ createdBy : chosen . ref . userId ,
580+ } )
583581
584- if ( result . length === 0 ) {
585- console . log ( ` [SKIP] BYOK already exists for provider "${ providerId } "` )
586- stats . skippedExisting ++
587- } else {
588- console . log ( ` [INSERT] BYOK for provider "${ providerId } ": ${ maskKey ( chosen . key ) } ` )
589- stats . inserted ++
590- }
582+ existingBYOKProviders . add ( providerId )
583+ console . log ( ` [INSERT] BYOK for provider "${ providerId } ": ${ maskKey ( chosen . key ) } ` )
584+ stats . inserted ++
591585 } catch ( error ) {
592586 console . error ( ` [ERROR] Failed to insert BYOK for provider "${ providerId } ":` , error )
593587 stats . errors ++
0 commit comments