[CONTAINS MIGRATION] User Invitations: Expiration and track accepted_…

…by_user_id (#11694)

airbyte-bootloader/src/test/java/io/airbyte/bootloader/BootloaderTest.java

@@ -97,7 +97,7 @@ class BootloaderTest {
 
   // ⚠️ This line should change with every new migration to show that you meant to make a new
   // migration to the prod database
-  private static final String CURRENT_CONFIGS_MIGRATION_VERSION = "0.50.41.010";
+  private static final String CURRENT_CONFIGS_MIGRATION_VERSION = "0.50.41.011";
   private static final String CURRENT_JOBS_MIGRATION_VERSION = "0.50.4.003";
   private static final String CDK_VERSION = "1.2.3";
 

airbyte-commons-server/src/main/java/io/airbyte/commons/server/errors/ConflictException.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.server.errors;
+
+import io.micronaut.http.HttpStatus;
+
+/**
+ * Exception when a request conflicts with the current state of the server. For example, trying to
+ * accept an invitation that was already accepted.
+ */
+public class ConflictException extends KnownException {
+
+  public ConflictException(final String message) {
+    super(message);
+  }
+
+  public ConflictException(final String message, final Throwable cause) {
+    super(message, cause);
+  }
+
+  @Override
+  public int getHttpCode() {
+    return HttpStatus.CONFLICT.getCode();
+  }
+
+}

airbyte-config/config-models/src/main/resources/types/InvitationStatus.yaml

@@ -9,3 +9,4 @@ enum:
   - accepted
   - cancelled
   - declined
+  - expired

airbyte-config/config-models/src/main/resources/types/UserInvitation.yaml

@@ -13,6 +13,7 @@ required:
   - scopeType
   - permissionType
   - status
+  - expiresAt
 additionalProperties: true
 properties:
   id:
@@ -29,6 +30,10 @@ properties:
     description: Email address of the user who is being invited
     type: string
     format: email
+  acceptedByUserId:
+    description: ID of the user who accepted the invitation
+    type: string
+    format: uuid
   scopeId:
     description: ID of the workspace/organization that the user is being invited to
     type: string
@@ -50,3 +55,7 @@ properties:
     description: last updated timestamp of the invitation
     type: integer
     format: int64
+  expiresAt:
+    description: Timestamp at which the invitation will expire
+    type: integer
+    format: int64

airbyte-data/src/main/kotlin/io/airbyte/data/repositories/entities/UserInvitation.kt

@@ -3,8 +3,10 @@ package io.airbyte.data.repositories.entities
 import io.airbyte.db.instance.configs.jooq.generated.enums.InvitationStatus
 import io.airbyte.db.instance.configs.jooq.generated.enums.PermissionType
 import io.airbyte.db.instance.configs.jooq.generated.enums.ScopeType
+import io.micronaut.core.annotation.Nullable
 import io.micronaut.data.annotation.AutoPopulated
 import io.micronaut.data.annotation.DateCreated
+import io.micronaut.data.annotation.DateUpdated
 import io.micronaut.data.annotation.Id
 import io.micronaut.data.annotation.MappedEntity
 import io.micronaut.data.annotation.TypeDef
@@ -19,6 +21,8 @@ data class UserInvitation(
   var inviteCode: String,
   var inviterUserId: UUID,
   var invitedEmail: String,
+  @Nullable
+  var acceptedByUserId: UUID? = null,
   var scopeId: UUID,
   @field:TypeDef(type = DataType.OBJECT)
   var scopeType: ScopeType,
@@ -28,6 +32,7 @@ data class UserInvitation(
   var status: InvitationStatus,
   @DateCreated
   var createdAt: java.time.OffsetDateTime? = null,
-  @DateCreated
+  @DateUpdated
   var updatedAt: java.time.OffsetDateTime? = null,
+  var expiresAt: java.time.OffsetDateTime,
 )

airbyte-data/src/main/kotlin/io/airbyte/data/services/UserInvitationService.kt

@@ -29,9 +29,10 @@ interface UserInvitationService {
   /**
    * Accept a user invitation and create resulting permission record.
    */
+  @Throws(InvitationStatusUnexpectedException::class)
   fun acceptUserInvitation(
     inviteCode: String,
-    invitedUserId: UUID,
+    acceptingUserId: UUID,
   ): UserInvitation
 
   /**
@@ -47,3 +48,9 @@ interface UserInvitationService {
    */
   fun cancelUserInvitation(inviteCode: String): UserInvitation
 }
+
+/**
+ * Exception thrown when an operation on an invitation cannot be performed because it has an
+ * unexpected status. For instance, trying to accept an invitation that is not pending.
+ */
+class InvitationStatusUnexpectedException(message: String) : Exception(message)

airbyte-data/src/main/kotlin/io/airbyte/data/services/impls/data/UserInvitationServiceDataImpl.kt

@@ -1,19 +1,21 @@
 package io.airbyte.data.services.impls.data
 
 import io.airbyte.config.ConfigSchema
-import io.airbyte.config.InvitationStatus
-import io.airbyte.config.Permission
 import io.airbyte.config.ScopeType
 import io.airbyte.config.UserInvitation
 import io.airbyte.data.exceptions.ConfigNotFoundException
 import io.airbyte.data.repositories.PermissionRepository
 import io.airbyte.data.repositories.UserInvitationRepository
+import io.airbyte.data.repositories.entities.Permission
+import io.airbyte.data.services.InvitationStatusUnexpectedException
 import io.airbyte.data.services.UserInvitationService
 import io.airbyte.data.services.impls.data.mappers.EntityInvitationStatus
+import io.airbyte.data.services.impls.data.mappers.EntityScopeType
 import io.airbyte.data.services.impls.data.mappers.toConfigModel
 import io.airbyte.data.services.impls.data.mappers.toEntity
 import io.micronaut.transaction.annotation.Transactional
 import jakarta.inject.Singleton
+import java.time.OffsetDateTime
 import java.util.UUID
 
 @Singleton
@@ -45,34 +47,45 @@ open class UserInvitationServiceDataImpl(
   @Transactional("config")
   override fun acceptUserInvitation(
     inviteCode: String,
-    invitedUserId: UUID,
+    acceptingUserId: UUID,
   ): UserInvitation {
     // fetch the invitation by code
     val invitation =
       userInvitationRepository.findByInviteCode(inviteCode).orElseThrow {
         ConfigNotFoundException(ConfigSchema.USER_INVITATION, inviteCode)
-      }.toConfigModel()
+      }
+
+    // mark the invitation status as expired if expiresAt is in the past
+    if (invitation.expiresAt.isBefore(OffsetDateTime.now())) {
+      invitation.status = EntityInvitationStatus.expired
+      userInvitationRepository.update(invitation)
+    }
 
-    if (invitation.status != InvitationStatus.PENDING) {
-      throw IllegalStateException("Invitation status is not pending: ${invitation.status}")
+    // throw an exception if the invitation is not pending. Note that this will also
+    // catch the case where the invitation is expired.
+    if (invitation.status != EntityInvitationStatus.pending) {
+      throw InvitationStatusUnexpectedException(
+        "Expected invitation for ScopeType: ${invitation.scopeType} and ScopeId: ${invitation.scopeId} to " +
+          "be PENDING, but instead it had Status: ${invitation.status}",
+      )
     }
 
     // create a new permission record according to the invitation
-    val permission =
-      Permission().apply {
-        userId = invitedUserId
-        permissionType = invitation.permissionType
-        when (invitation.scopeType) {
-          ScopeType.ORGANIZATION -> organizationId = invitation.scopeId
-          ScopeType.WORKSPACE -> workspaceId = invitation.scopeId
-          else -> throw IllegalStateException("Unknown scope type: ${invitation.scopeType}")
-        }
+    Permission(
+      id = UUID.randomUUID(),
+      userId = acceptingUserId,
+      permissionType = invitation.permissionType,
+    ).apply {
+      when (invitation.scopeType) {
+        EntityScopeType.organization -> organizationId = invitation.scopeId
+        EntityScopeType.workspace -> workspaceId = invitation.scopeId
       }
-    permissionRepository.save(permission.toEntity())
+    }.let { permissionRepository.save(it) }
 
-    // update the invitation status to accepted
-    invitation.status = InvitationStatus.ACCEPTED
-    val updatedInvitation = userInvitationRepository.update(invitation.toEntity())
+    // mark the invitation as accepted
+    invitation.status = EntityInvitationStatus.accepted
+    invitation.acceptedByUserId = acceptingUserId
+    val updatedInvitation = userInvitationRepository.update(invitation)
 
     return updatedInvitation.toConfigModel()
   }

airbyte-data/src/main/kotlin/io/airbyte/data/services/impls/data/mappers/UserInvitationMapper.kt

@@ -32,6 +32,7 @@ fun EntityInvitationStatus.toConfigModel(): ModelInvitationStatus {
     EntityInvitationStatus.accepted -> ModelInvitationStatus.ACCEPTED
     EntityInvitationStatus.cancelled -> ModelInvitationStatus.CANCELLED
     EntityInvitationStatus.declined -> ModelInvitationStatus.DECLINED
+    EntityInvitationStatus.expired -> ModelInvitationStatus.EXPIRED
   }
 }
 
@@ -41,6 +42,7 @@ fun ModelInvitationStatus.toEntity(): EntityInvitationStatus {
     ModelInvitationStatus.ACCEPTED -> EntityInvitationStatus.accepted
     ModelInvitationStatus.CANCELLED -> EntityInvitationStatus.cancelled
     ModelInvitationStatus.DECLINED -> EntityInvitationStatus.declined
+    ModelInvitationStatus.EXPIRED -> EntityInvitationStatus.expired
   }
 }
 
@@ -50,12 +52,14 @@ fun EntityUserInvitation.toConfigModel(): ModelUserInvitation {
     .withInviteCode(this.inviteCode)
     .withInviterUserId(this.inviterUserId)
     .withInvitedEmail(this.invitedEmail)
+    .withAcceptedByUserId(this.acceptedByUserId)
     .withScopeId(this.scopeId)
     .withScopeType(this.scopeType.toConfigModel())
     .withPermissionType(this.permissionType.toConfigModel())
     .withStatus(this.status.toConfigModel())
     .withCreatedAt(this.createdAt?.toEpochSecond())
     .withUpdatedAt(this.updatedAt?.toEpochSecond())
+    .withExpiresAt(this.expiresAt.toEpochSecond())
 }
 
 fun ModelUserInvitation.toEntity(): EntityUserInvitation {
@@ -64,11 +68,13 @@ fun ModelUserInvitation.toEntity(): EntityUserInvitation {
     inviteCode = this.inviteCode,
     inviterUserId = this.inviterUserId,
     invitedEmail = this.invitedEmail,
+    acceptedByUserId = this.acceptedByUserId,
     scopeId = this.scopeId,
     scopeType = this.scopeType.toEntity(),
     permissionType = this.permissionType.toEntity(),
     status = this.status.toEntity(),
     createdAt = this.createdAt?.let { OffsetDateTime.ofInstant(Instant.ofEpochSecond(it), ZoneOffset.UTC) },
     updatedAt = this.updatedAt?.let { OffsetDateTime.ofInstant(Instant.ofEpochSecond(it), ZoneOffset.UTC) },
+    expiresAt = OffsetDateTime.ofInstant(Instant.ofEpochSecond(this.expiresAt), ZoneOffset.UTC),
   )
 }

airbyte-data/src/test/kotlin/io/airbyte/data/repositories/UserInvitationRepositoryTest.kt

@@ -11,12 +11,15 @@ import io.micronaut.test.extensions.junit5.annotation.MicronautTest
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.BeforeAll
 import org.junit.jupiter.api.Test
+import java.time.OffsetDateTime
+import java.time.ZoneOffset
 import java.util.UUID
 
 @MicronautTest
 internal class UserInvitationRepositoryTest : AbstractConfigRepositoryTest<UserInvitationRepository>(UserInvitationRepository::class) {
   companion object {
     const val INVITE_CODE = "some-code"
+    val EXPIRES_AT = OffsetDateTime.now(ZoneOffset.UTC).plusDays(7).truncatedTo(java.time.temporal.ChronoUnit.SECONDS)
 
     val userInvitation =
       UserInvitation(
@@ -27,6 +30,7 @@ internal class UserInvitationRepositoryTest : AbstractConfigRepositoryTest<UserI
         scopeType = ScopeType.workspace,
         permissionType = PermissionType.workspace_admin,
         status = InvitationStatus.pending,
+        expiresAt = EXPIRES_AT,
       )
 
     @BeforeAll

airbyte-data/src/test/kotlin/io/airbyte/data/services/impls/data/UserInvitationServiceDataImplTest.kt

@@ -6,6 +6,7 @@ import io.airbyte.data.repositories.PermissionRepository
 import io.airbyte.data.repositories.UserInvitationRepository
 import io.airbyte.data.repositories.entities.Permission
 import io.airbyte.data.repositories.entities.UserInvitation
+import io.airbyte.data.services.InvitationStatusUnexpectedException
 import io.airbyte.data.services.impls.data.mappers.EntityInvitationStatus
 import io.airbyte.data.services.impls.data.mappers.EntityPermissionType
 import io.airbyte.data.services.impls.data.mappers.EntityScopeType
@@ -42,6 +43,7 @@ internal class UserInvitationServiceDataImplTest {
       status = EntityInvitationStatus.pending,
       createdAt = OffsetDateTime.now(ZoneOffset.UTC).truncatedTo(java.time.temporal.ChronoUnit.SECONDS),
       updatedAt = OffsetDateTime.now(ZoneOffset.UTC).truncatedTo(java.time.temporal.ChronoUnit.SECONDS),
+      expiresAt = OffsetDateTime.now(ZoneOffset.UTC).plusDays(7).truncatedTo(java.time.temporal.ChronoUnit.SECONDS),
     )
 
   @BeforeEach
@@ -81,7 +83,11 @@ internal class UserInvitationServiceDataImplTest {
   @Test
   fun `test accept user invitation`() {
     val invitedUserId = UUID.randomUUID()
-    val expectedUpdatedInvitation = invitation.copy(status = EntityInvitationStatus.accepted)
+    val expectedUpdatedInvitation =
+      invitation.copy(
+        status = EntityInvitationStatus.accepted,
+        acceptedByUserId = invitedUserId,
+      )
 
     every { userInvitationRepository.findByInviteCode(invitation.inviteCode) } returns Optional.of(invitation)
     every { userInvitationRepository.update(expectedUpdatedInvitation) } returns expectedUpdatedInvitation
@@ -117,11 +123,29 @@ internal class UserInvitationServiceDataImplTest {
 
     every { userInvitationRepository.findByInviteCode(invitation.inviteCode) } returns Optional.of(invitation)
 
-    assertThrows<IllegalStateException> { userInvitationService.acceptUserInvitation(invitation.inviteCode, invitedUserId) }
+    assertThrows<InvitationStatusUnexpectedException> { userInvitationService.acceptUserInvitation(invitation.inviteCode, invitedUserId) }
 
     verify(exactly = 0) { userInvitationRepository.update(any()) }
   }
 
+  @Test
+  fun `test accept user invitation fails if expired`() {
+    val invitedUserId = UUID.randomUUID()
+    val expiredInvitation =
+      invitation.copy(
+        status = EntityInvitationStatus.pending,
+        expiresAt = OffsetDateTime.now(ZoneOffset.UTC).minusDays(1),
+      )
+    val expectedUpdatedInvitation = expiredInvitation.copy(status = EntityInvitationStatus.expired)
+
+    every { userInvitationRepository.findByInviteCode(expiredInvitation.inviteCode) } returns Optional.of(expiredInvitation)
+    every { userInvitationRepository.update(expectedUpdatedInvitation) } returns expectedUpdatedInvitation
+
+    assertThrows<InvitationStatusUnexpectedException> { userInvitationService.acceptUserInvitation(expiredInvitation.inviteCode, invitedUserId) }
+
+    verify { userInvitationRepository.update(expectedUpdatedInvitation) }
+  }
+
   @Test
   fun `test get pending invitations`() {
     val workspaceId = UUID.randomUUID()