feat: connector config entitlement for rejected records storage (#16799)

Author: pedroslopez

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

@@ -10,6 +10,7 @@ import io.airbyte.config.ActorType
 import io.airbyte.featureflag.AllowConfigTemplateEndpoints
 import io.airbyte.featureflag.DestinationDefinition
 import io.airbyte.featureflag.FeatureFlagClient
+import io.airbyte.featureflag.LicenseAllowDestinationObjectStorageConfig
 import io.airbyte.featureflag.LicenseAllowEnterpriseConnector
 import io.airbyte.featureflag.Multi
 import io.airbyte.featureflag.Organization
@@ -28,6 +29,8 @@ interface EntitlementProvider {
   ): Map<UUID, Boolean>
 
   fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean
+
+  fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean
 }
 
 /**
@@ -42,6 +45,8 @@ class DefaultEntitlementProvider : EntitlementProvider {
   ): Map<UUID, Boolean> = actorDefinitionIds.associateWith { _ -> false }
 
   override fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean = false
+
+  override fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean = false
 }
 
 /**
@@ -68,6 +73,8 @@ class EnterpriseEntitlementProvider(
   }
 
   override fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean = activeLicense.license?.isEmbedded ?: false
+
+  override fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean = true
 }
 
 /**
@@ -104,4 +111,7 @@ class CloudEntitlementProvider(
 
   override fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean =
     featureFlagClient.boolVariation(AllowConfigTemplateEndpoints, Organization(organizationId))
+
+  override fun hasDestinationObjectStorageEntitlement(organizationId: UUID): Boolean =
+    featureFlagClient.boolVariation(flag = LicenseAllowDestinationObjectStorageConfig, Organization(organizationId))
 }

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

@@ -4,7 +4,10 @@
 
 package io.airbyte.commons.entitlements
 
+import io.airbyte.api.problems.model.generated.ProblemLicenseEntitlementData
+import io.airbyte.api.problems.throwable.generated.LicenseEntitlementProblem
 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.config.ActorType
@@ -19,7 +22,24 @@ class EntitlementService(
   fun checkEntitlement(
     organizationId: UUID,
     entitlement: Entitlement,
-  ): EntitlementResult = entitlementClient.checkEntitlement(organizationId, entitlement)
+  ): EntitlementResult =
+    when (entitlement) {
+      // TODO: Remove once we've migrated the entitlement to Stigg
+      DestinationObjectStorageEntitlement -> hasDestinationObjectStorageEntitlement(organizationId)
+      else -> entitlementClient.checkEntitlement(organizationId, entitlement)
+    }
+
+  fun ensureEntitled(
+    organizationId: UUID,
+    entitlement: Entitlement,
+  ) {
+    if (!checkEntitlement(organizationId, entitlement).isEntitled) {
+      throw LicenseEntitlementProblem(
+        ProblemLicenseEntitlementData()
+          .entitlement(entitlement.featureId),
+      )
+    }
+  }
 
   fun getEntitlements(organizationId: UUID): List<EntitlementResult> = entitlementClient.getEntitlements(organizationId)
 
@@ -42,5 +62,11 @@ class EntitlementService(
     return clientResults.toMutableMap().apply { putAll(providerResults) }
   }
 
+  private fun hasDestinationObjectStorageEntitlement(organizationId: UUID): EntitlementResult =
+    EntitlementResult(
+      isEntitled = entitlementProvider.hasDestinationObjectStorageEntitlement(organizationId),
+      featureId = DestinationObjectStorageEntitlement.featureId,
+    )
+
   internal fun hasConfigTemplateEntitlements(organizationId: UUID): Boolean = entitlementProvider.hasConfigTemplateEntitlements(organizationId)
 }

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

@@ -12,11 +12,16 @@ object PlatformSubOneHourSyncFrequency : FeatureEntitlement(
   featureId = "feature-platform-sub-one-hour-sync-frequency",
 )
 
+object DestinationObjectStorageEntitlement : FeatureEntitlement(
+  featureId = "feature-destination-object-storage",
+)
+
 object Entitlements {
   private val ALL =
     listOf(
       PlatformLlmSyncJobFailureExplanation,
       PlatformSubOneHourSyncFrequency,
+      DestinationObjectStorageEntitlement,
     )
 
   private val byId = ALL.associateBy { it.featureId }

airbyte-commons-entitlements/src/test/kotlin/io/airbyte/commons/entitlements/EntitlementProviderTest.kt

@@ -9,6 +9,7 @@ import io.airbyte.commons.license.AirbyteLicense
 import io.airbyte.config.ActorType
 import io.airbyte.featureflag.AllowConfigTemplateEndpoints
 import io.airbyte.featureflag.DestinationDefinition
+import io.airbyte.featureflag.LicenseAllowDestinationObjectStorageConfig
 import io.airbyte.featureflag.LicenseAllowEnterpriseConnector
 import io.airbyte.featureflag.Multi
 import io.airbyte.featureflag.Organization
@@ -51,6 +52,13 @@ class EntitlementProviderTest {
       val res = entitlementProvider.hasConfigTemplateEntitlements(organizationId)
       assertFalse(res)
     }
+
+    @Test
+    fun `test hasDestinationObjectStorageEntitlement`() {
+      val organizationId = UUID.randomUUID()
+      val res = entitlementProvider.hasDestinationObjectStorageEntitlement(organizationId)
+      assertFalse(res)
+    }
   }
 
   @Nested
@@ -93,6 +101,13 @@ class EntitlementProviderTest {
       val res = entitlementProvider.hasConfigTemplateEntitlements(organizationId)
       assertEquals(res, license.isEmbedded)
     }
+
+    @Test
+    fun `test hasDestinationObjectStorageEntitlement always returns true`() {
+      val organizationId = UUID.randomUUID()
+      val res = entitlementProvider.hasDestinationObjectStorageEntitlement(organizationId)
+      assertEquals(true, res)
+    }
   }
 
   @Nested
@@ -132,6 +147,16 @@ class EntitlementProviderTest {
       assertEquals(res, isEntitled)
     }
 
+    @ParameterizedTest
+    @ValueSource(booleans = [true, false])
+    fun `test hasDestinationObjectStorageEntitlement returns value from feature flag`(isEntitled: Boolean) {
+      val organizationId = UUID.randomUUID()
+      every { featureFlagClient.boolVariation(LicenseAllowDestinationObjectStorageConfig, Organization(organizationId)) } returns isEntitled
+
+      val res = entitlementProvider.hasDestinationObjectStorageEntitlement(organizationId)
+      assertEquals(res, isEntitled)
+    }
+
     private fun mockEntitledEnterpriseConnector(
       actorType: ActorType,
       organizationId: UUID,

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

@@ -29,6 +29,9 @@
 import io.airbyte.data.services.DestinationService;
 import io.airbyte.data.services.OAuthService;
 import io.airbyte.data.services.SourceService;
+import io.airbyte.domain.models.EntitledConnectorSpec;
+import io.airbyte.domain.services.entitlements.ConnectorConfigEntitlementService;
+import io.airbyte.persistence.job.WorkspaceHelper;
 import io.airbyte.validation.json.JsonValidationException;
 import jakarta.inject.Singleton;
 import java.io.IOException;
@@ -50,17 +53,23 @@ public class ConnectorDefinitionSpecificationHandler {
   private final SourceService sourceService;
   private final DestinationService destinationService;
   private final OAuthService oAuthService;
+  private final ConnectorConfigEntitlementService connectorConfigEntitlementService;
+  private final WorkspaceHelper workspaceHelper;
 
   public ConnectorDefinitionSpecificationHandler(final ActorDefinitionVersionHelper actorDefinitionVersionHelper,
                                                  final JobConverter jobConverter,
                                                  final SourceService sourceService,
                                                  final DestinationService destinationService,
-                                                 final OAuthService oauthService) {
+                                                 final WorkspaceHelper workspaceHelper,
+                                                 final OAuthService oauthService,
+                                                 final ConnectorConfigEntitlementService connectorConfigEntitlementService) {
     this.actorDefinitionVersionHelper = actorDefinitionVersionHelper;
     this.jobConverter = jobConverter;
     this.sourceService = sourceService;
     this.destinationService = destinationService;
     this.oAuthService = oauthService;
+    this.connectorConfigEntitlementService = connectorConfigEntitlementService;
+    this.workspaceHelper = workspaceHelper;
   }
 
   /**
@@ -78,9 +87,9 @@ public SourceDefinitionSpecificationRead getSpecificationForSourceId(final Sourc
     final StandardSourceDefinition sourceDefinition = sourceService.getStandardSourceDefinition(source.getSourceDefinitionId());
     final ActorDefinitionVersion sourceVersion =
         actorDefinitionVersionHelper.getSourceVersion(sourceDefinition, source.getWorkspaceId(), sourceIdRequestBody.getSourceId());
-    final io.airbyte.protocol.models.v0.ConnectorSpecification spec = sourceVersion.getSpec();
-
-    return getSourceSpecificationRead(sourceDefinition, spec, source.getWorkspaceId());
+    final EntitledConnectorSpec entitledConnectorSpec =
+        connectorConfigEntitlementService.getEntitledConnectorSpec(source.getWorkspaceId(), sourceVersion);
+    return getSourceSpecificationRead(sourceDefinition, entitledConnectorSpec, source.getWorkspaceId());
   }
 
   /**
@@ -98,9 +107,10 @@ public SourceDefinitionSpecificationRead getSourceDefinitionSpecification(final
     final StandardSourceDefinition source = sourceService.getStandardSourceDefinition(sourceDefinitionId);
     final ActorDefinitionVersion sourceVersion =
         actorDefinitionVersionHelper.getSourceVersion(source, sourceDefinitionIdWithWorkspaceId.getWorkspaceId());
-    final io.airbyte.protocol.models.v0.ConnectorSpecification spec = sourceVersion.getSpec();
-
-    return getSourceSpecificationRead(source, spec, sourceDefinitionIdWithWorkspaceId.getWorkspaceId());
+    final UUID organizationId = workspaceHelper.getOrganizationForWorkspace(sourceDefinitionIdWithWorkspaceId.getWorkspaceId());
+    final EntitledConnectorSpec entitledConnectorSpec =
+        connectorConfigEntitlementService.getEntitledConnectorSpec(organizationId, sourceVersion);
+    return getSourceSpecificationRead(source, entitledConnectorSpec, sourceDefinitionIdWithWorkspaceId.getWorkspaceId());
   }
 
   /**
@@ -120,8 +130,11 @@ public DestinationDefinitionSpecificationRead getSpecificationForDestinationId(f
     final ActorDefinitionVersion destinationVersion =
         actorDefinitionVersionHelper.getDestinationVersion(destinationDefinition, destination.getWorkspaceId(),
             destinationIdRequestBody.getDestinationId());
-    final io.airbyte.protocol.models.v0.ConnectorSpecification spec = destinationVersion.getSpec();
-    return getDestinationSpecificationRead(destinationDefinition, spec, destinationVersion.getSupportsRefreshes(), destination.getWorkspaceId());
+    final UUID organizationId = workspaceHelper.getOrganizationForWorkspace(destination.getWorkspaceId());
+    final EntitledConnectorSpec entitledConnectorSpec =
+        connectorConfigEntitlementService.getEntitledConnectorSpec(organizationId, destinationVersion);
+    return getDestinationSpecificationRead(destinationDefinition, entitledConnectorSpec, destinationVersion.getSupportsRefreshes(),
+        destination.getWorkspaceId());
   }
 
   /**
@@ -141,17 +154,19 @@ public DestinationDefinitionSpecificationRead getDestinationSpecification(final
     final StandardDestinationDefinition destination = destinationService.getStandardDestinationDefinition(destinationDefinitionId);
     final ActorDefinitionVersion destinationVersion =
         actorDefinitionVersionHelper.getDestinationVersion(destination, destinationDefinitionIdWithWorkspaceId.getWorkspaceId());
-    final io.airbyte.protocol.models.v0.ConnectorSpecification spec = destinationVersion.getSpec();
-
-    return getDestinationSpecificationRead(destination, spec, destinationVersion.getSupportsRefreshes(),
+    final UUID organizationId = workspaceHelper.getOrganizationForWorkspace(destinationDefinitionIdWithWorkspaceId.getWorkspaceId());
+    final EntitledConnectorSpec entitledConnectorSpec =
+        connectorConfigEntitlementService.getEntitledConnectorSpec(organizationId, destinationVersion);
+    return getDestinationSpecificationRead(destination, entitledConnectorSpec, destinationVersion.getSupportsRefreshes(),
         destinationDefinitionIdWithWorkspaceId.getWorkspaceId());
   }
 
   @VisibleForTesting
   SourceDefinitionSpecificationRead getSourceSpecificationRead(final StandardSourceDefinition sourceDefinition,
-                                                               final io.airbyte.protocol.models.v0.ConnectorSpecification spec,
+                                                               final EntitledConnectorSpec entitledConnectorSpec,
                                                                final UUID workspaceId)
       throws IOException {
+    final io.airbyte.protocol.models.v0.ConnectorSpecification spec = entitledConnectorSpec.getSpec();
     final SourceDefinitionSpecificationRead specRead = new SourceDefinitionSpecificationRead()
         .jobInfo(jobConverter.getSynchronousJobRead(SynchronousJobMetadata.mock(JobConfig.ConfigType.GET_SPEC)))
         .connectionSpecification(spec.getConnectionSpecification())
@@ -174,10 +189,11 @@ SourceDefinitionSpecificationRead getSourceSpecificationRead(final StandardSourc
 
   @VisibleForTesting
   DestinationDefinitionSpecificationRead getDestinationSpecificationRead(final StandardDestinationDefinition destinationDefinition,
-                                                                         final io.airbyte.protocol.models.v0.ConnectorSpecification spec,
+                                                                         final EntitledConnectorSpec entitledConnectorSpec,
                                                                          final boolean supportsRefreshes,
                                                                          final UUID workspaceId)
       throws IOException {
+    final io.airbyte.protocol.models.v0.ConnectorSpecification spec = entitledConnectorSpec.getSpec();
     final DestinationDefinitionSpecificationRead specRead = new DestinationDefinitionSpecificationRead()
         .jobInfo(jobConverter.getSynchronousJobRead(SynchronousJobMetadata.mock(JobConfig.ConfigType.GET_SPEC)))
         .supportedDestinationSyncModes(getFinalDestinationSyncModes(spec.getSupportedDestinationSyncModes(), supportsRefreshes))