Persist geography updates (#18501)

* persist geography in connection update * use ApiPojoConverters * add test coverage for persisting geography update * persist geography column for workspace * use US instead of AUTO in mock data so that tests don't pass due to the database-level default column value * use ApiPojoConverters instead of Enum.convertTo in workspace handler * format
airbytehq · Oct 26, 2022 · 5d4b564 · 5d4b564
1 parent 9e612b0
commit 5d4b564
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 5 deletions.
diff --git a/...ig-persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java b/...ig-persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java
@@ -768,6 +768,9 @@ private void writeStandardWorkspace(final List<StandardWorkspace> configs, final
             .set(WORKSPACE.NOTIFICATIONS, JSONB.valueOf(Jsons.serialize(standardWorkspace.getNotifications())))
             .set(WORKSPACE.FIRST_SYNC_COMPLETE, standardWorkspace.getFirstCompletedSync())
             .set(WORKSPACE.FEEDBACK_COMPLETE, standardWorkspace.getFeedbackDone())
+            .set(WORKSPACE.GEOGRAPHY, Enums.toEnum(
+                standardWorkspace.getDefaultGeography().value(),
+                io.airbyte.db.instance.configs.jooq.generated.enums.GeographyType.class).orElseThrow())
             .set(WORKSPACE.UPDATED_AT, timestamp)
             .set(WORKSPACE.WEBHOOK_OPERATION_CONFIGS, standardWorkspace.getWebhookOperationConfigs() == null ? null
                 : JSONB.valueOf(Jsons.serialize(standardWorkspace.getWebhookOperationConfigs())))
@@ -791,6 +794,9 @@ private void writeStandardWorkspace(final List<StandardWorkspace> configs, final
             .set(WORKSPACE.FEEDBACK_COMPLETE, standardWorkspace.getFeedbackDone())
             .set(WORKSPACE.CREATED_AT, timestamp)
             .set(WORKSPACE.UPDATED_AT, timestamp)
+            .set(WORKSPACE.GEOGRAPHY, Enums.toEnum(
+                standardWorkspace.getDefaultGeography().value(),
+                io.airbyte.db.instance.configs.jooq.generated.enums.GeographyType.class).orElseThrow())
             .set(WORKSPACE.WEBHOOK_OPERATION_CONFIGS, standardWorkspace.getWebhookOperationConfigs() == null ? null
                 : JSONB.valueOf(Jsons.serialize(standardWorkspace.getWebhookOperationConfigs())))
             .execute();

diff --git a/airbyte-config/config-persistence/src/test/java/io/airbyte/config/persistence/MockData.java b/airbyte-config/config-persistence/src/test/java/io/airbyte/config/persistence/MockData.java
@@ -161,7 +161,7 @@ public static List<StandardWorkspace> standardWorkspaces() {
         .withNotifications(Collections.singletonList(notification))
         .withFirstCompletedSync(true)
         .withFeedbackDone(true)
-        .withDefaultGeography(Geography.AUTO)
+        .withDefaultGeography(Geography.US)
         .withWebhookOperationConfigs(Jsons.jsonNode(
             new WebhookOperationConfigs().withWebhookConfigs(List.of(new WebhookConfig().withId(WEBHOOK_CONFIG_ID).withName("name")))));
 

diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/ConnectionsHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/ConnectionsHandler.java
@@ -339,6 +339,10 @@ private static void applyPatchToStandardSync(final StandardSync sync, final Conn
     if (patch.getResourceRequirements() != null) {
       sync.setResourceRequirements(ApiPojoConverters.resourceRequirementsToInternal(patch.getResourceRequirements()));
     }
+
+    if (patch.getGeography() != null) {
+      sync.setGeography(ApiPojoConverters.toPersistenceGeography(patch.getGeography()));
+    }
   }
 
   private void validateConnectionPatch(final WorkspaceHelper workspaceHelper, final StandardSync persistedSync, final ConnectionUpdate patch) {

diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/WorkspacesHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/WorkspacesHandler.java
@@ -30,6 +30,7 @@
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.config.persistence.SecretsRepositoryWriter;
 import io.airbyte.notification.NotificationClient;
+import io.airbyte.server.converters.ApiPojoConverters;
 import io.airbyte.server.converters.NotificationConverter;
 import io.airbyte.server.converters.WorkspaceWebhookConfigsConverter;
 import io.airbyte.server.errors.IdNotFoundKnownException;
@@ -301,8 +302,7 @@ private void applyPatchToStandardWorkspace(final StandardWorkspace workspace, fi
       workspace.setNotifications(NotificationConverter.toConfigList(workspacePatch.getNotifications()));
     }
     if (workspacePatch.getDefaultGeography() != null) {
-      workspace.setDefaultGeography(
-          Enums.convertTo(workspacePatch.getDefaultGeography(), io.airbyte.config.Geography.class));
+      workspace.setDefaultGeography(ApiPojoConverters.toPersistenceGeography(workspacePatch.getDefaultGeography()));
     }
     if (workspacePatch.getWebhookConfigs() != null) {
       workspace.setWebhookOperationConfigs(WorkspaceWebhookConfigsConverter.toPersistenceWrite(workspacePatch.getWebhookConfigs(), uuidSupplier));

diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/ConnectionsHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/ConnectionsHandlerTest.java
@@ -586,7 +586,8 @@ void testUpdateConnectionPatchingSeveralFieldsAndReplaceAStream() throws JsonVal
             .syncCatalog(catalogForUpdate)
             .resourceRequirements(resourceRequirements)
             .sourceCatalogId(newSourceCatalogId)
-            .operationIds(List.of(operationId, otherOperationId));
+            .operationIds(List.of(operationId, otherOperationId))
+            .geography(io.airbyte.api.model.generated.Geography.EU);
 
         final ConfiguredAirbyteCatalog expectedPersistedCatalog = ConnectionHelpers.generateBasicConfiguredAirbyteCatalog();
         expectedPersistedCatalog.getStreams().get(0).getStream().withName(AZKABAN_USERS);
@@ -600,7 +601,8 @@ void testUpdateConnectionPatchingSeveralFieldsAndReplaceAStream() throws JsonVal
             .withCatalog(expectedPersistedCatalog)
             .withResourceRequirements(ApiPojoConverters.resourceRequirementsToInternal(resourceRequirements))
             .withSourceCatalogId(newSourceCatalogId)
-            .withOperationIds(List.of(operationId, otherOperationId));
+            .withOperationIds(List.of(operationId, otherOperationId))
+            .withGeography(Geography.EU);
 
         when(configRepository.getStandardSync(standardSync.getConnectionId())).thenReturn(standardSync);