From a63d6777336819dffd9bd8953e8497a548100cfa Mon Sep 17 00:00:00 2001 From: pmossman Date: Fri, 13 Jan 2023 11:01:39 -0800 Subject: [PATCH] add repo method to determine if workspace has an alpha or beta connector --- .../config/persistence/ConfigRepository.java | 24 ++++++ .../persistence/WorkspacePersistenceTest.java | 81 +++++++++++++++++++ 2 files changed, 105 insertions(+) diff --git a/airbyte-config/config-persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java b/airbyte-config/config-persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java index d398b851759da..6865993a0126f 100644 --- a/airbyte-config/config-persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java +++ b/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. * diff --git a/airbyte-config/config-persistence/src/test/java/io/airbyte/config/persistence/WorkspacePersistenceTest.java b/airbyte-config/config-persistence/src/test/java/io/airbyte/config/persistence/WorkspacePersistenceTest.java index 8077d3fdcbc60..0d740aa12d921 100644 --- a/airbyte-config/config-persistence/src/test/java/io/airbyte/config/persistence/WorkspacePersistenceTest.java +++ b/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)); + } + }