Actually use bootstrap handler (#11907)

airbyte-api/src/main/openapi/config.yaml

@@ -6164,6 +6164,7 @@ components:
       type: object
       required:
         - name
+        - organizationId
       properties:
         email:
           type: string
@@ -6197,6 +6198,7 @@ components:
       required:
         - id
         - name
+        - organizationId
       properties:
         id:
           type: string

airbyte-commons-server/src/main/java/io/airbyte/commons/server/handlers/UserHandler.java

@@ -25,6 +25,7 @@
 import io.airbyte.api.model.generated.UserUpdate;
 import io.airbyte.api.model.generated.UserWithPermissionInfoRead;
 import io.airbyte.api.model.generated.UserWithPermissionInfoReadList;
+import io.airbyte.api.model.generated.WorkspaceCreateWithId;
 import io.airbyte.api.model.generated.WorkspaceIdRequestBody;
 import io.airbyte.api.model.generated.WorkspaceRead;
 import io.airbyte.api.model.generated.WorkspaceReadList;
@@ -37,6 +38,7 @@
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.server.errors.ConflictException;
 import io.airbyte.commons.server.errors.OperationNotAllowedException;
+import io.airbyte.commons.server.handlers.helpers.WorkspaceHelpersKt;
 import io.airbyte.commons.server.support.UserAuthenticationResolver;
 import io.airbyte.config.ConfigSchema;
 import io.airbyte.config.Organization;
@@ -84,6 +86,7 @@ public class UserHandler {
 
   private final UserAuthenticationResolver userAuthenticationResolver;
   private final Optional<InitialUserConfiguration> initialUserConfiguration;
+  private final ResourceBootstrapHandlerInterface resourceBootstrapHandler;
 
   @VisibleForTesting
   public UserHandler(
@@ -95,7 +98,8 @@ public UserHandler(
                      final WorkspacesHandler workspacesHandler,
                      @Named("uuidGenerator") final Supplier<UUID> uuidGenerator,
                      final UserAuthenticationResolver userAuthenticationResolver,
-                     final Optional<InitialUserConfiguration> initialUserConfiguration) {
+                     final Optional<InitialUserConfiguration> initialUserConfiguration,
+                     final ResourceBootstrapHandlerInterface resourceBootstrapHandler) {
     this.uuidGenerator = uuidGenerator;
     this.userPersistence = userPersistence;
     this.organizationPersistence = organizationPersistence;
@@ -105,6 +109,7 @@ public UserHandler(
     this.permissionHandler = permissionHandler;
     this.userAuthenticationResolver = userAuthenticationResolver;
     this.initialUserConfiguration = initialUserConfiguration;
+    this.resourceBootstrapHandler = resourceBootstrapHandler;
   }
 
   /**
@@ -386,9 +391,11 @@ private void handleUserPermissionsAndWorkspace(final UserRead createdUser)
     createInstanceAdminPermissionIfInitialUser(createdUser);
     final Optional<Organization> ssoOrg = getSsoOrganizationIfExists();
     if (ssoOrg.isPresent()) {
+      // SSO users will have some additional logic but will ultimately call createDefaultWorkspaceForUser
       handleSsoUser(createdUser, ssoOrg.get());
     } else {
-      handleNonSsoUser(createdUser);
+      // non-SSO users will just create a default workspace
+      createDefaultWorkspaceForUser(createdUser, Optional.empty());
     }
   }
 
@@ -416,28 +423,42 @@ private void handleSsoUser(final UserRead user, final Organization organization)
         new ListWorkspacesInOrganizationRequestBody().organizationId(organization.getOrganizationId()));
 
     if (orgWorkspaces.getWorkspaces().isEmpty()) {
-      final WorkspaceRead defaultWorkspace = createDefaultWorkspaceForUser(user, Optional.of(organization));
-      createPermissionForUserAndWorkspace(user.getUserId(), defaultWorkspace.getWorkspaceId(), PermissionType.WORKSPACE_ADMIN);
+      // Now calls bootstrap which includes all permissions and updates userRead.
+      createDefaultWorkspaceForUser(user, Optional.of(organization));
     }
   }
 
-  private void handleNonSsoUser(final UserRead user) throws JsonValidationException, ConfigNotFoundException, IOException {
-    final WorkspaceRead defaultWorkspace = createDefaultWorkspaceForUser(user, Optional.empty());
-    createPermissionForUserAndWorkspace(user.getUserId(), defaultWorkspace.getWorkspaceId(), PermissionType.WORKSPACE_ADMIN);
-  }
-
-  private WorkspaceRead createDefaultWorkspaceForUser(final UserRead createdUser, final Optional<Organization> organization)
+  protected void createDefaultWorkspaceForUser(final UserRead user, final Optional<Organization> organization)
       throws JsonValidationException, IOException, ConfigNotFoundException {
 
-    final WorkspaceRead defaultWorkspace = workspacesHandler.createDefaultWorkspaceForUser(createdUser, organization);
+    // Only do this if the user doesn't already have a default workspace.
+    if (user.getDefaultWorkspaceId() != null) {
+      return;
+    }
+
+    // Logic stolen from workspaceHandler.createDefaultWorkspaceForUser
+    final String companyName = user.getCompanyName();
+    final String email = user.getEmail();
+    final Boolean news = user.getNews();
+    // otherwise, create a default workspace for this user
+    final WorkspaceCreateWithId workspaceCreate = new WorkspaceCreateWithId()
+        .name(WorkspaceHelpersKt.getDefaultWorkspaceName(organization, companyName, email))
+        .organizationId(organization.map(Organization::getOrganizationId).orElse(null))
+        .email(email)
+        .news(news)
+        .anonymousDataCollection(false)
+        .securityUpdates(false)
+        .displaySetupWizard(true)
+        .id(uuidGenerator.get());
+
+    final WorkspaceRead defaultWorkspace = resourceBootstrapHandler.bootStrapWorkspaceForCurrentUser(workspaceCreate);
 
     // set default workspace id in User table
     final UserUpdate userUpdateDefaultWorkspace = new UserUpdate()
-        .userId(createdUser.getUserId())
+        .userId(user.getUserId())
         .defaultWorkspaceId(defaultWorkspace.getWorkspaceId());
     updateUser(userUpdateDefaultWorkspace);
 
-    return defaultWorkspace;
   }
 
   private Optional<Organization> getSsoOrganizationIfExists() throws IOException {

airbyte-commons-server/src/main/java/io/airbyte/commons/server/handlers/WorkspacesHandler.java

@@ -39,8 +39,10 @@
 import io.airbyte.commons.server.converters.NotificationSettingsConverter;
 import io.airbyte.commons.server.converters.WorkspaceConverter;
 import io.airbyte.commons.server.converters.WorkspaceWebhookConfigsConverter;
+import io.airbyte.commons.server.errors.BadObjectSchemaKnownException;
 import io.airbyte.commons.server.errors.InternalServerKnownException;
 import io.airbyte.commons.server.errors.ValueConflictKnownException;
+import io.airbyte.commons.server.handlers.helpers.WorkspaceHelpersKt;
 import io.airbyte.config.Organization;
 import io.airbyte.config.StandardWorkspace;
 import io.airbyte.config.persistence.ConfigNotFoundException;
@@ -137,6 +139,13 @@ public WorkspaceRead createWorkspace(final WorkspaceCreate workspaceCreate)
   public WorkspaceRead createWorkspaceIfNotExist(final WorkspaceCreateWithId workspaceCreateWithId)
       throws JsonValidationException, IOException, ValueConflictKnownException, ConfigNotFoundException {
 
+    // We expect that the caller is specifying the workspace ID.
+    // Since this code is currently only called by OSS, it's enforced in the public API and the UI
+    // currently.
+    if (workspaceCreateWithId.getOrganizationId() == null) {
+      throw new BadObjectSchemaKnownException("Workspace missing org ID.");
+    }
+
     final String email = workspaceCreateWithId.getEmail();
     final Boolean anonymousDataCollection = workspaceCreateWithId.getAnonymousDataCollection();
     final Boolean news = workspaceCreateWithId.getNews();
@@ -187,7 +196,7 @@ public WorkspaceRead createDefaultWorkspaceForUser(final UserRead user, final Op
     final Boolean news = user.getNews();
     // otherwise, create a default workspace for this user
     final WorkspaceCreate workspaceCreate = new WorkspaceCreate()
-        .name(getDefaultWorkspaceName(organization, companyName, email))
+        .name(WorkspaceHelpersKt.getDefaultWorkspaceName(organization, companyName, email))
         .organizationId(organization.map(Organization::getOrganizationId).orElse(null))
         .email(email)
         .news(news)
@@ -197,24 +206,6 @@ public WorkspaceRead createDefaultWorkspaceForUser(final UserRead user, final Op
     return createWorkspace(workspaceCreate);
   }
 
-  private String getDefaultWorkspaceName(final Optional<Organization> organization, final String companyName, final String email) {
-    String defaultWorkspaceName = "";
-    if (organization.isPresent()) {
-      // use organization name as default workspace name
-      defaultWorkspaceName = organization.get().getName().trim();
-    }
-    // if organization name is not available or empty, use user's company name (note: this is an
-    // optional field)
-    if (defaultWorkspaceName.isEmpty() && companyName != null) {
-      defaultWorkspaceName = companyName.trim();
-    }
-    // if company name is still empty, use user's email (note: this is a required field)
-    if (defaultWorkspaceName.isEmpty()) {
-      defaultWorkspaceName = email;
-    }
-    return defaultWorkspaceName;
-  }
-
   public void deleteWorkspace(final WorkspaceIdRequestBody workspaceIdRequestBody)
       throws JsonValidationException, IOException, ConfigNotFoundException {
     // get existing implementation

airbyte-commons-server/src/main/kotlin/io/airbyte/commons/server/handlers/ResourceBootstrapHandler.kt

@@ -36,15 +36,15 @@ open class ResourceBootstrapHandler(
   private val permissionService: PermissionService,
   private val currentUserService: CurrentUserService,
   private val apiAuthorizationHelper: ApiAuthorizationHelper,
-) {
+) : ResourceBootstrapHandlerInterface {
   companion object {
     val LOGGER = LoggerFactory.getLogger(ResourceBootstrapHandler::class.java)
   }
 
   /**
    * This is for bootstrapping a workspace and all the necessary links (organization) and permissions (workspace & organization).
    */
-  fun bootStrapWorkspaceForCurrentUser(workspaceCreateWithId: WorkspaceCreateWithId): WorkspaceRead {
+  override fun bootStrapWorkspaceForCurrentUser(workspaceCreateWithId: WorkspaceCreateWithId): WorkspaceRead {
     val user = currentUserService.getCurrentUser()
     // The organization to use to set up the new workspace
     val organization =
@@ -81,7 +81,7 @@ open class ResourceBootstrapHandler(
     return WorkspaceConverter.domainToApiModel(standardWorkspace)
   }
 
-  private fun findOrCreateOrganizationAndPermission(user: User): Organization {
+  public fun findOrCreateOrganizationAndPermission(user: User): Organization {
     findExistingOrganization(user)?.let { return it }
 
     val organization =

airbyte-commons-server/src/main/kotlin/io/airbyte/commons/server/handlers/ResourceBootstrapHandlerInterface.kt

@@ -0,0 +1,8 @@
+package io.airbyte.commons.server.handlers
+
+import io.airbyte.api.model.generated.WorkspaceCreateWithId
+import io.airbyte.api.model.generated.WorkspaceRead
+
+interface ResourceBootstrapHandlerInterface {
+  fun bootStrapWorkspaceForCurrentUser(workspaceCreateWithId: WorkspaceCreateWithId): WorkspaceRead
+}

airbyte-commons-server/src/main/kotlin/io/airbyte/commons/server/handlers/helpers/WorkspaceHelpers.kt

@@ -11,9 +11,12 @@ import io.airbyte.commons.server.converters.WorkspaceWebhookConfigsConverter
 import io.airbyte.config.Geography
 import io.airbyte.config.Organization
 import io.airbyte.config.StandardWorkspace
+import java.util.Optional
 import java.util.UUID
 import java.util.function.Supplier
 
+// These helpers exist so that we can get some of the utility of working with workspaces but without needing to inject WorkspacesHandler
+
 fun buildStandardWorkspace(
   workspaceCreateWithId: WorkspaceCreateWithId,
   organization: Organization,
@@ -40,8 +43,7 @@ fun buildStandardWorkspace(
   val notificationSettings: NotificationSettings = patchNotificationSettingsWithDefaultValue(workspaceCreateWithId)
 
   return StandardWorkspace().apply {
-    this.workspaceId = uuidSupplier.get()
-    this.workspaceId = workspaceCreateWithId.id
+    this.workspaceId = workspaceCreateWithId.id ?: uuidSupplier.get()
     this.customerId = uuidSupplier.get() // "customer_id" should be deprecated
     this.name = workspaceCreateWithId.name
     this.slug = uuidSupplier.get().toString()
@@ -78,3 +80,25 @@ private fun patchNotificationSettingsWithDefaultValue(workspaceCreateWithId: Wor
       workspaceCreateWithId.notificationSettings?.sendOnBreakingChangeSyncsDisabled ?: defaultNotificationType
   }
 }
+
+fun getDefaultWorkspaceName(
+  organization: Optional<Organization>,
+  companyName: String?,
+  email: String,
+): String {
+  var defaultWorkspaceName = ""
+  if (organization.isPresent) {
+    // use organization name as default workspace name
+    defaultWorkspaceName = organization.get().name.trim()
+  }
+  // if organization name is not available or empty, use user's company name (note: this is an
+  // optional field)
+  if (defaultWorkspaceName.isEmpty() && companyName != null) {
+    defaultWorkspaceName = companyName.trim()
+  }
+  // if company name is still empty, use user's email (note: this is a required field)
+  if (defaultWorkspaceName.isEmpty()) {
+    defaultWorkspaceName = email
+  }
+  return defaultWorkspaceName
+}

airbyte-commons-server/src/test/java/io/airbyte/commons/server/handlers/UserHandlerTest.java

@@ -10,7 +10,6 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -100,6 +99,7 @@ class UserHandlerTest {
       .withAuthProvider(AuthProvider.GOOGLE_IDENTITY_PLATFORM)
       .withStatus(Status.INVITED)
       .withName(USER_NAME);
+  private ResourceBootstrapHandler resourceBootstrapHandler;
 
   @BeforeEach
   void setUp() {
@@ -113,10 +113,11 @@ void setUp() {
     uuidSupplier = mock(Supplier.class);
     jwtUserAuthenticationResolver = mock(JwtUserAuthenticationResolver.class);
     initialUserConfiguration = mock(InitialUserConfiguration.class);
+    resourceBootstrapHandler = mock(ResourceBootstrapHandler.class);
 
     userHandler =
         new UserHandler(userPersistence, permissionPersistence, permissionService, organizationPersistence, permissionHandler, workspacesHandler,
-            uuidSupplier, jwtUserAuthenticationResolver, Optional.of(initialUserConfiguration));
+            uuidSupplier, jwtUserAuthenticationResolver, Optional.of(initialUserConfiguration), resourceBootstrapHandler);
   }
 
   @Test
@@ -328,7 +329,7 @@ void setUp() throws IOException, JsonValidationException, ConfigNotFoundExceptio
         when(jwtUserAuthenticationResolver.resolveUser(NEW_AUTH_USER_ID)).thenReturn(newUser);
         when(uuidSupplier.get()).thenReturn(NEW_USER_ID);
         when(userPersistence.getUser(NEW_USER_ID)).thenReturn(Optional.of(newUser));
-        when(workspacesHandler.createDefaultWorkspaceForUser(any(), any())).thenReturn(defaultWorkspace);
+        when(resourceBootstrapHandler.bootStrapWorkspaceForCurrentUser(any())).thenReturn(defaultWorkspace);
       }
 
       @ParameterizedTest
@@ -357,7 +358,7 @@ void testNewUserCreation(final AuthProvider authProvider,
           // happens in Cloud)
           userHandler = new UserHandler(userPersistence, permissionPersistence, permissionService, organizationPersistence, permissionHandler,
               workspacesHandler,
-              uuidSupplier, jwtUserAuthenticationResolver, Optional.empty());
+              uuidSupplier, jwtUserAuthenticationResolver, Optional.empty(), resourceBootstrapHandler);
         }
 
         if (isFirstOrgUser) {
@@ -376,6 +377,9 @@ void testNewUserCreation(final AuthProvider authProvider,
           when(workspacesHandler.listWorkspacesInOrganization(
               new ListWorkspacesInOrganizationRequestBody().organizationId(ORGANIZATION.getOrganizationId()))).thenReturn(
                   new WorkspaceReadList().workspaces(List.of(defaultWorkspace)));
+          if (newUser.getDefaultWorkspaceId() == null) {
+            newUser.setDefaultWorkspaceId(defaultWorkspace.getWorkspaceId());
+          }
         } else {
           when(workspacesHandler.listWorkspacesInOrganization(any())).thenReturn(new WorkspaceReadList().workspaces(List.of()));
         }
@@ -395,7 +399,7 @@ void testNewUserCreation(final AuthProvider authProvider,
         verifyUserRead(userRead, apiAuthProvider);
         verifyInstanceAdminPermissionCreation(initialUserEmail, initialUserPresent);
         verifyOrganizationPermissionCreation(ssoRealm, isFirstOrgUser);
-        verifyDefaultWorkspaceCreation(ssoRealm, isDefaultWorkspaceForOrgPresent, userPersistenceInOrder);
+        verifyDefaultWorkspaceCreation(isDefaultWorkspaceForOrgPresent, userPersistenceInOrder);
       }
 
       private void verifyCreatedUser(final AuthProvider expectedAuthProvider, final InOrder inOrder) throws IOException {
@@ -405,38 +409,19 @@ private void verifyCreatedUser(final AuthProvider expectedAuthProvider, final In
             && user.getAuthProvider().equals(expectedAuthProvider)));
       }
 
-      private void verifyDefaultWorkspaceCreation(final String ssoRealm, final Boolean isDefaultWorkspaceForOrgPresent, final InOrder inOrder)
-          throws IOException, JsonValidationException, ConfigNotFoundException {
-        boolean workspaceCreated = false;
-
-        if (ssoRealm == null) {
-          // always create a default workspace for non-SSO users
-          verify(workspacesHandler).createDefaultWorkspaceForUser(
-              argThat(user -> user.getUserId().equals(NEW_USER_ID)),
-              eq(Optional.empty()));
-          workspaceCreated = true;
-
-        } else {
-          if (!isDefaultWorkspaceForOrgPresent) {
-            // create a default workspace for the org if one doesn't yet exist
-            verify(workspacesHandler).createDefaultWorkspaceForUser(
-                argThat(user -> user.getUserId().equals(NEW_USER_ID)),
-                argThat(org -> org.orElseThrow().getOrganizationId().equals(ORGANIZATION.getOrganizationId())));
-            workspaceCreated = true;
-
-          } else {
-            // never create an additional workspace for the org if one already exists.
-            verify(workspacesHandler, never()).createDefaultWorkspaceForUser(any(), any());
-          }
-        }
-        if (workspaceCreated) {
+      private void verifyDefaultWorkspaceCreation(final Boolean isDefaultWorkspaceForOrgPresent, final InOrder inOrder)
+          throws IOException {
+        // No need to deal with other vars because SSO users and first org users etc. are all directed
+        // through the same codepath now.
+        if (!isDefaultWorkspaceForOrgPresent) {
+          // create a default workspace for the org if one doesn't yet exist
+          verify(resourceBootstrapHandler).bootStrapWorkspaceForCurrentUser(any());
           // if a workspace was created, verify that the user's defaultWorkspaceId was updated
           // and that a workspaceAdmin permission was created for them.
           inOrder.verify(userPersistence).writeUser(argThat(user -> user.getDefaultWorkspaceId().equals(WORKSPACE_ID)));
-          verify(permissionHandler).createPermission(new PermissionCreate()
-              .permissionType(io.airbyte.api.model.generated.PermissionType.WORKSPACE_ADMIN)
-              .workspaceId(WORKSPACE_ID)
-              .userId(NEW_USER_ID));
+        } else {
+          // never create an additional workspace for the org if one already exists.
+          verify(resourceBootstrapHandler, never()).bootStrapWorkspaceForCurrentUser(any());
         }
       }