feat: add sso config feature flag (#16961)

Author: cjkenned

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

@@ -20415,7 +20415,7 @@ components:
       properties:
         organizationId:
           type: string
-          format: UUID
+          format: uuid
         companyIdentifier:
           type: string
           description: Used to name the keycloak realm
@@ -20437,7 +20437,7 @@ components:
       properties:
         organizationId:
           type: string
-          format: UUID
+          format: uuid
         companyIdentifier:
           type: string
           description: Matches the keycloak realm to be removed

airbyte-commons-entitlements/src/main/kotlin/io/airbyte/commons/entitlements/EntitlementProvider.kt

@@ -9,6 +9,7 @@ import io.airbyte.commons.license.annotation.RequiresAirbyteProEnabled
 import io.airbyte.config.ActorType
 import io.airbyte.featureflag.AllowConfigTemplateEndpoints
 import io.airbyte.featureflag.DestinationDefinition
+import io.airbyte.featureflag.EnableSsoConfigUpdate
 import io.airbyte.featureflag.FeatureFlagClient
 import io.airbyte.featureflag.LicenseAllowDestinationObjectStorageConfig
 import io.airbyte.featureflag.LicenseAllowEnterpriseConnector
@@ -31,6 +32,8 @@ interface EntitlementProvider {
   fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean
 
   fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean
+
+  fun hasSsoConfigUpdateEntitlement(organizationId: UUID): Boolean
 }
 
 /**
@@ -47,6 +50,8 @@ class DefaultEntitlementProvider : EntitlementProvider {
   override fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean = false
 
   override fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean = false
+
+  override fun hasSsoConfigUpdateEntitlement(organizationId: UUID): Boolean = false
 }
 
 /**
@@ -75,6 +80,8 @@ class EnterpriseEntitlementProvider(
   override fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean = activeLicense.license?.isEmbedded ?: false
 
   override fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean = true
+
+  override fun hasSsoConfigUpdateEntitlement(organizationId: UUID): Boolean = false
 }
 
 /**
@@ -114,4 +121,7 @@ class CloudEntitlementProvider(
 
   override fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean =
     featureFlagClient.boolVariation(flag = LicenseAllowDestinationObjectStorageConfig, Organization(organizationId))
+
+  override fun hasSsoConfigUpdateEntitlement(organizationId: UUID): Boolean =
+    featureFlagClient.boolVariation(EnableSsoConfigUpdate, Organization(organizationId))
 }

airbyte-commons-entitlements/src/main/kotlin/io/airbyte/commons/entitlements/EntitlementService.kt

@@ -10,6 +10,7 @@ import io.airbyte.commons.entitlements.models.ConnectorEntitlement
 import io.airbyte.commons.entitlements.models.DestinationObjectStorageEntitlement
 import io.airbyte.commons.entitlements.models.Entitlement
 import io.airbyte.commons.entitlements.models.EntitlementResult
+import io.airbyte.commons.entitlements.models.SsoConfigUpdateEntitlement
 import io.airbyte.config.ActorType
 import jakarta.inject.Singleton
 import java.util.UUID
@@ -26,6 +27,7 @@ class EntitlementService(
     when (entitlement) {
       // TODO: Remove once we've migrated the entitlement to Stigg
       DestinationObjectStorageEntitlement -> hasDestinationObjectStorageEntitlement(organizationId)
+      SsoConfigUpdateEntitlement -> hasSsoConfigUpdateEntitlement(organizationId)
       else -> entitlementClient.checkEntitlement(organizationId, entitlement)
     }
 
@@ -68,5 +70,11 @@ class EntitlementService(
       featureId = DestinationObjectStorageEntitlement.featureId,
     )
 
+  private fun hasSsoConfigUpdateEntitlement(organizationId: UUID): EntitlementResult =
+    EntitlementResult(
+      isEntitled = entitlementProvider.hasSsoConfigUpdateEntitlement(organizationId),
+      featureId = SsoConfigUpdateEntitlement.featureId,
+    )
+
   internal fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean = entitlementProvider.hasConfigTemplateEntitlements(organizationId)
 }

airbyte-commons-entitlements/src/main/kotlin/io/airbyte/commons/entitlements/models/EntitlementDefinitions.kt

@@ -16,12 +16,17 @@ object DestinationObjectStorageEntitlement : FeatureEntitlement(
   featureId = "feature-destination-object-storage",
 )
 
+object SsoConfigUpdateEntitlement : FeatureEntitlement(
+  featureId = "feature-platform-sso-config-update",
+)
+
 object Entitlements {
   private val ALL =
     listOf(
       PlatformLlmSyncJobFailureExplanation,
       PlatformSubOneHourSyncFrequency,
       DestinationObjectStorageEntitlement,
+      SsoConfigUpdateEntitlement,
     )
 
   private val byId = ALL.associateBy { it.featureId }

airbyte-featureflag/src/main/kotlin/FlagDefinitions.kt

@@ -217,3 +217,6 @@ object EnableDestinationCatalogValidation : Temporary<Boolean>(key = "platform.e
 object LicenseAllowDestinationObjectStorageConfig : Permanent<Boolean>(key = "license.allow-destination-object-storage-config", default = false)
 
 object UseSonarServer : Temporary<Boolean>(key = "embedded.useSonarServer", default = false)
+
+// this uses the webapp naming conventions, as the flag is used in both the frontend and platform.
+object EnableSsoConfigUpdate : Permanent<Boolean>(key = "featureService.ALLOW_UPDATE_SSO_CONFIG", default = false)

airbyte-server/src/main/kotlin/io/airbyte/server/apis/controllers/SSOConfigApiController.kt

@@ -13,6 +13,8 @@ import io.airbyte.api.server.generated.models.UpdateSSOCredentialsRequestBody
 import io.airbyte.commons.annotation.AuditLogging
 import io.airbyte.commons.annotation.AuditLoggingProvider
 import io.airbyte.commons.auth.roles.AuthRoleConstants
+import io.airbyte.commons.entitlements.EntitlementService
+import io.airbyte.commons.entitlements.models.SsoConfigUpdateEntitlement
 import io.airbyte.commons.server.scheduling.AirbyteTaskExecutors
 import io.airbyte.domain.models.SsoConfig
 import io.airbyte.domain.models.SsoKeycloakIdpCredentials
@@ -21,15 +23,17 @@ import io.airbyte.server.apis.execute
 import io.micronaut.http.annotation.Controller
 import io.micronaut.scheduling.annotation.ExecuteOn
 import io.micronaut.security.annotation.Secured
-import java.util.UUID
 
 @Controller
 open class SSOConfigApiController(
   private val ssoConfigDomainService: SsoConfigDomainService,
+  private val entitlementService: EntitlementService,
 ) : SsoConfigApi {
   @Secured(AuthRoleConstants.ORGANIZATION_ADMIN)
   @ExecuteOn(AirbyteTaskExecutors.IO)
   override fun getSsoConfig(getSSOConfigRequestBody: GetSSOConfigRequestBody): SSOConfigRead {
+    entitlementService.ensureEntitled(getSSOConfigRequestBody.organizationId, SsoConfigUpdateEntitlement)
+
     val ssoConfig = ssoConfigDomainService.retrieveSsoConfig(getSSOConfigRequestBody.organizationId)
     return SSOConfigRead(
       organizationId = getSSOConfigRequestBody.organizationId,
@@ -44,10 +48,12 @@ open class SSOConfigApiController(
   @ExecuteOn(AirbyteTaskExecutors.IO)
   @AuditLogging(provider = AuditLoggingProvider.BASIC)
   override fun createSsoConfig(createSSOConfigRequestBody: CreateSSOConfigRequestBody) {
+    entitlementService.ensureEntitled(createSSOConfigRequestBody.organizationId, SsoConfigUpdateEntitlement)
+
     execute<Any?> {
       ssoConfigDomainService.createAndStoreSsoConfig(
         SsoConfig(
-          UUID.fromString(createSSOConfigRequestBody.organizationId),
+          createSSOConfigRequestBody.organizationId,
           createSSOConfigRequestBody.companyIdentifier,
           createSSOConfigRequestBody.clientId,
           createSSOConfigRequestBody.clientSecret,
@@ -63,9 +69,11 @@ open class SSOConfigApiController(
   @ExecuteOn(AirbyteTaskExecutors.IO)
   @AuditLogging(provider = AuditLoggingProvider.BASIC)
   override fun deleteSsoConfig(deleteSSOConfigRequestBody: DeleteSSOConfigRequestBody) {
+    entitlementService.ensureEntitled(deleteSSOConfigRequestBody.organizationId, SsoConfigUpdateEntitlement)
+
     execute<Any?> {
       ssoConfigDomainService.deleteSsoConfig(
-        UUID.fromString(deleteSSOConfigRequestBody.organizationId),
+        deleteSSOConfigRequestBody.organizationId,
         deleteSSOConfigRequestBody.companyIdentifier,
       )
     }
@@ -75,6 +83,8 @@ open class SSOConfigApiController(
   @ExecuteOn(AirbyteTaskExecutors.IO)
   @AuditLogging(provider = AuditLoggingProvider.BASIC)
   override fun updateSsoCredentials(updateSSOCredentialsRequestBody: UpdateSSOCredentialsRequestBody) {
+    entitlementService.ensureEntitled(updateSSOCredentialsRequestBody.organizationId, SsoConfigUpdateEntitlement)
+
     execute<Any?> {
       ssoConfigDomainService.updateClientCredentials(
         SsoKeycloakIdpCredentials(

airbyte-server/src/test/kotlin/io/airbyte/server/apis/controllers/SsoConfigApiControllerTest.kt

@@ -0,0 +1,115 @@
+/*
+ * Copyright (c) 2020-2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.server.apis.controllers
+
+import io.airbyte.api.server.generated.models.CreateSSOConfigRequestBody
+import io.airbyte.api.server.generated.models.DeleteSSOConfigRequestBody
+import io.airbyte.api.server.generated.models.GetSSOConfigRequestBody
+import io.airbyte.api.server.generated.models.UpdateSSOCredentialsRequestBody
+import io.airbyte.commons.entitlements.EntitlementService
+import io.airbyte.commons.entitlements.models.SsoConfigUpdateEntitlement
+import io.airbyte.domain.models.SsoConfigRetrieval
+import io.airbyte.domain.services.sso.SsoConfigDomainService
+import io.mockk.Runs
+import io.mockk.every
+import io.mockk.just
+import io.mockk.mockk
+import io.mockk.verify
+import org.junit.jupiter.api.Test
+import java.util.UUID
+
+class SsoConfigApiControllerTest {
+  companion object {
+    private val ssoConfigDomainService = mockk<SsoConfigDomainService>()
+    private val entitlementService = mockk<EntitlementService>()
+    private val ssoConfigController =
+      SSOConfigApiController(
+        ssoConfigDomainService,
+        entitlementService,
+      )
+  }
+
+  @Test
+  fun `getSsoConfig returns the config`() {
+    val orgId = UUID.randomUUID()
+    every { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) } just Runs
+    every { ssoConfigDomainService.retrieveSsoConfig(orgId) } returns
+      SsoConfigRetrieval(
+        companyIdentifier = "id",
+        clientId = "client-id",
+        clientSecret = "client-secret",
+        emailDomains = listOf("domain"),
+      )
+
+    val result =
+      ssoConfigController.getSsoConfig(
+        GetSSOConfigRequestBody(orgId),
+      )
+
+    verify(exactly = 1) { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) }
+
+    assert(result.organizationId == orgId)
+    assert(result.companyIdentifier == "id")
+    assert(result.clientId == "client-id")
+    assert(result.clientSecret == "client-secret")
+    assert(result.emailDomains == listOf("domain"))
+  }
+
+  @Test
+  fun `createSsoConfig creates a new config`() {
+    val orgId = UUID.randomUUID()
+    every { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) } just Runs
+    every { ssoConfigDomainService.createAndStoreSsoConfig(any()) } just Runs
+
+    ssoConfigController.createSsoConfig(
+      CreateSSOConfigRequestBody(
+        organizationId = orgId,
+        companyIdentifier = "id",
+        clientId = "client-id",
+        clientSecret = "client-secret",
+        discoveryUrl = "https://www.airbyte.io",
+        emailDomain = "domain",
+      ),
+    )
+
+    verify(exactly = 1) { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) }
+    verify(exactly = 1) { ssoConfigDomainService.createAndStoreSsoConfig(any()) }
+  }
+
+  @Test
+  fun `deleteSsoConfig removes the existing config`() {
+    val orgId = UUID.randomUUID()
+    every { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) } just Runs
+    every { ssoConfigDomainService.deleteSsoConfig(orgId, any()) } just Runs
+
+    ssoConfigController.deleteSsoConfig(
+      DeleteSSOConfigRequestBody(
+        organizationId = orgId,
+        companyIdentifier = "id",
+      ),
+    )
+
+    verify(exactly = 1) { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) }
+    verify(exactly = 1) { ssoConfigDomainService.deleteSsoConfig(orgId, any()) }
+  }
+
+  @Test
+  fun `updateSsoConfig updates a new config`() {
+    val orgId = UUID.randomUUID()
+    every { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) } just Runs
+    every { ssoConfigDomainService.updateClientCredentials(any()) } just Runs
+
+    ssoConfigController.updateSsoCredentials(
+      UpdateSSOCredentialsRequestBody(
+        organizationId = orgId,
+        clientId = "client-id",
+        clientSecret = "client-secret",
+      ),
+    )
+
+    verify(exactly = 1) { entitlementService.ensureEntitled(orgId, SsoConfigUpdateEntitlement) }
+    verify(exactly = 1) { ssoConfigDomainService.updateClientCredentials(any()) }
+  }
+}