Add query to determine whether a workspace has an alpha or beta connector

airbyte-config/config-persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java

@@ -1517,6 +1517,30 @@ public Geography getGeographyForConnection(final UUID connectionId) throws IOExc
         .fetchOneInto(Geography.class);
   }
 
+  /**
+   * Specialized query for efficiently determining eligibility for the Free Connector Program. If a
+   * workspace has at least one Alpha or Beta connector, users of that workspace will be prompted to
+   * sign up for the program. This check is performed on nearly every page load so the query needs to
+   * be as efficient as possible.
+   *
+   * @param workspaceId ID of the workspace to check connectors for
+   * @return boolean indicating if an alpha or beta connector exists within the workspace
+   */
+  public boolean getWorkspaceHasAlphaOrBetaConnector(final UUID workspaceId) throws IOException {
+    final Condition releaseStageAlphaOrBeta = ACTOR_DEFINITION.RELEASE_STAGE.eq(ReleaseStage.alpha)
+        .or(ACTOR_DEFINITION.RELEASE_STAGE.eq(ReleaseStage.beta));
+
+    final Integer countResult = database.query(ctx -> ctx.selectCount()
+        .from(ACTOR)
+        .join(ACTOR_DEFINITION).on(ACTOR_DEFINITION.ID.eq(ACTOR.ACTOR_DEFINITION_ID))
+        .where(ACTOR.WORKSPACE_ID.eq(workspaceId))
+        .and(ACTOR.TOMBSTONE.notEqual(true))
+        .and(releaseStageAlphaOrBeta))
+        .fetchOneInto(Integer.class);
+
+    return countResult > 0;
+  }
+
   /**
    * Deletes all records with given id. If it deletes anything, returns true. Otherwise, false.
    *

airbyte-config/config-persistence/src/test/java/io/airbyte/config/persistence/WorkspacePersistenceTest.java

@@ -11,10 +11,15 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.lang.MoreBooleans;
+import io.airbyte.config.DestinationConnection;
 import io.airbyte.config.Geography;
 import io.airbyte.config.SourceConnection;
+import io.airbyte.config.StandardDestinationDefinition;
+import io.airbyte.config.StandardSourceDefinition;
+import io.airbyte.config.StandardSourceDefinition.ReleaseStage;
 import io.airbyte.config.StandardSync;
 import io.airbyte.config.StandardWorkspace;
 import io.airbyte.db.ExceptionWrappingDatabase;
@@ -25,11 +30,17 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
 
 @SuppressWarnings({"PMD.LongVariable", "PMD.AvoidInstantiatingObjectsInLoops"})
 class WorkspacePersistenceTest extends BaseConfigDatabaseTest {
 
   private static final UUID WORKSPACE_ID = UUID.randomUUID();
+  private static final UUID SOURCE_DEFINITION_ID = UUID.randomUUID();
+  private static final UUID SOURCE_ID = UUID.randomUUID();
+  private static final UUID DESTINATION_DEFINITION_ID = UUID.randomUUID();
+  private static final UUID DESTINATION_ID = UUID.randomUUID();
+  private static final JsonNode CONFIG = Jsons.jsonNode(ImmutableMap.of("key-a", "value-a"));
 
   private ConfigRepository configRepository;
 
@@ -72,6 +83,57 @@ private static StandardWorkspace createBaseStandardWorkspace() {
         .withDefaultGeography(Geography.AUTO);
   }
 
+  private static SourceConnection createBaseSource() {
+    return new SourceConnection()
+        .withSourceId(SOURCE_ID)
+        .withSourceDefinitionId(SOURCE_DEFINITION_ID)
+        .withName("source-a")
+        .withTombstone(false)
+        .withConfiguration(CONFIG)
+        .withWorkspaceId(WORKSPACE_ID);
+  }
+
+  private static DestinationConnection createBaseDestination() {
+    return new DestinationConnection()
+        .withDestinationId(DESTINATION_ID)
+        .withDestinationDefinitionId(DESTINATION_DEFINITION_ID)
+        .withName("destination-a")
+        .withTombstone(false)
+        .withConfiguration(CONFIG)
+        .withWorkspaceId(WORKSPACE_ID);
+  }
+
+  private static StandardSourceDefinition createSourceDefinition(final ReleaseStage releaseStage) {
+    return new StandardSourceDefinition()
+        .withSourceDefinitionId(SOURCE_DEFINITION_ID)
+        .withTombstone(false)
+        .withName("source-definition-a")
+        .withDockerRepository("dockerhub")
+        .withDockerImageTag("some-tag")
+        .withReleaseStage(releaseStage);
+  }
+
+  private static StandardDestinationDefinition createDestinationDefinition(final StandardDestinationDefinition.ReleaseStage releaseStage) {
+    return new StandardDestinationDefinition()
+        .withDestinationDefinitionId(DESTINATION_DEFINITION_ID)
+        .withTombstone(false)
+        .withName("destination-definition-a")
+        .withDockerRepository("dockerhub")
+        .withDockerImageTag("some-tag")
+        .withReleaseStage(releaseStage);
+  }
+
+  private void persistConnectorsWithReleaseStages(
+                                                  final ReleaseStage sourceReleaseStage,
+                                                  final StandardDestinationDefinition.ReleaseStage destinationReleaseStage)
+      throws JsonValidationException, IOException {
+
+    configRepository.writeStandardSourceDefinition(createSourceDefinition(sourceReleaseStage));
+    configRepository.writeStandardDestinationDefinition(createDestinationDefinition(destinationReleaseStage));
+    configRepository.writeSourceConnectionNoSecrets(createBaseSource());
+    configRepository.writeDestinationConnectionNoSecrets(createBaseDestination());
+  }
+
   void assertReturnsWorkspace(final StandardWorkspace workspace) throws ConfigNotFoundException, IOException, JsonValidationException {
     configRepository.writeStandardWorkspaceNoSecrets(workspace);
 
@@ -125,4 +187,23 @@ void testUpdateFeedback() throws JsonValidationException, ConfigNotFoundExceptio
     assertTrue(configRepository.getStandardWorkspaceNoSecrets(workspace.getWorkspaceId(), false).getFeedbackDone());
   }
 
+  @Test
+  void testWorkspaceHasAlphaOrBetaConnector() throws JsonValidationException, IOException {
+    final StandardWorkspace workspace = createBaseStandardWorkspace();
+
+    configRepository.writeStandardWorkspaceNoSecrets(workspace);
+
+    persistConnectorsWithReleaseStages(ReleaseStage.GENERALLY_AVAILABLE, StandardDestinationDefinition.ReleaseStage.GENERALLY_AVAILABLE);
+    assertFalse(configRepository.getWorkspaceHasAlphaOrBetaConnector(WORKSPACE_ID));
+
+    persistConnectorsWithReleaseStages(ReleaseStage.ALPHA, StandardDestinationDefinition.ReleaseStage.GENERALLY_AVAILABLE);
+    assertTrue(configRepository.getWorkspaceHasAlphaOrBetaConnector(WORKSPACE_ID));
+
+    persistConnectorsWithReleaseStages(ReleaseStage.GENERALLY_AVAILABLE, StandardDestinationDefinition.ReleaseStage.BETA);
+    assertTrue(configRepository.getWorkspaceHasAlphaOrBetaConnector(WORKSPACE_ID));
+
+    persistConnectorsWithReleaseStages(ReleaseStage.CUSTOM, StandardDestinationDefinition.ReleaseStage.CUSTOM);
+    assertFalse(configRepository.getWorkspaceHasAlphaOrBetaConnector(WORKSPACE_ID));
+  }
+
 }