Add Workspace and Connection Geography Support to API (#17650)

* progress on adding geography throughout api

* fix workspace handler test

* more progress

* implement workspace defaulting and add/update more tests

* fix bootloader tests

* set defaultGeography in missing places

* add Geography column when reading Connection record from DB

* fix pmd

* add more comments/description

* format

* description

airbyte-analytics/src/test/java/io/airbyte/analytics/TrackingClientSingletonTest.java

@@ -12,6 +12,7 @@
 import io.airbyte.commons.version.AirbyteVersion;
 import io.airbyte.config.Configs;
 import io.airbyte.config.Configs.WorkerEnvironment;
+import io.airbyte.config.Geography;
 import io.airbyte.config.StandardWorkspace;
 import io.airbyte.config.persistence.ConfigNotFoundException;
 import io.airbyte.config.persistence.ConfigRepository;
@@ -111,7 +112,8 @@ void testGetTrackingIdentityNonAnonymous() throws JsonValidationException, IOExc
         .withEmail(EMAIL)
         .withAnonymousDataCollection(false)
         .withNews(true)
-        .withSecurityUpdates(true);
+        .withSecurityUpdates(true)
+        .withDefaultGeography(Geography.AUTO);
 
     when(configRepository.getStandardWorkspace(WORKSPACE_ID, true)).thenReturn(workspace);
 
@@ -129,7 +131,8 @@ void testGetTrackingIdentityAnonymous() throws JsonValidationException, IOExcept
         .withEmail("a@airbyte.io")
         .withAnonymousDataCollection(true)
         .withNews(true)
-        .withSecurityUpdates(true);
+        .withSecurityUpdates(true)
+        .withDefaultGeography(Geography.AUTO);
 
     when(configRepository.getStandardWorkspace(WORKSPACE_ID, true)).thenReturn(workspace);
 

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

@@ -2036,6 +2036,24 @@ paths:
           $ref: "#/components/responses/NotFoundResponse"
         "422":
           $ref: "#/components/responses/InvalidInputResponse"
+  /v1/web_backend/geographies/list:
+    post:
+      tags:
+        - web_backend
+      description: Returns all available geographies in which a data sync can run.
+      summary: |
+        Returns available geographies can be selected to run data syncs in a particular geography. 
+        The 'auto' entry indicates that the sync will be automatically assigned to a geography according 
+        to the platform default behavior. Entries other than 'auto' are two-letter country codes that 
+        follow the ISO 3166-1 alpha-2 standard.
+      operationId: webBackendListGeographies
+      responses:
+        "200":
+          description: Successful operation
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/WebBackendGeographiesListResult"
   /v1/jobs/list:
     post:
       tags:
@@ -2260,6 +2278,8 @@ components:
             $ref: "#/components/schemas/Notification"
         displaySetupWizard:
           type: boolean
+        defaultGeography:
+          $ref: "#/components/schemas/Geography"
     Notification:
       type: object
       required:
@@ -2362,6 +2382,8 @@ components:
           type: boolean
         feedbackDone:
           type: boolean
+        defaultGeography:
+          $ref: "#/components/schemas/Geography"
     WorkspaceUpdateName:
       type: object
       required:
@@ -2397,6 +2419,8 @@ components:
           type: array
           items:
             $ref: "#/components/schemas/Notification"
+        defaultGeography:
+          $ref: "#/components/schemas/Geography"
     WorkspaceGiveFeedback:
       type: object
       required:
@@ -2424,6 +2448,15 @@ components:
           type: boolean
         hasDestinations:
           type: boolean
+    WebBackendGeographiesListResult:
+      type: object
+      required:
+        - geographies
+      properties:
+        geographies:
+          type: array
+          items:
+            $ref: "#/components/schemas/Geography"
     # SLUG
     SlugRequestBody:
       type: object
@@ -2432,6 +2465,13 @@ components:
       properties:
         slug:
           type: string
+    # Geography
+    Geography:
+      type: string
+      enum:
+        - auto
+        - us
+        - eu
     # SourceDefinition
     SourceDefinitionId:
       type: string
@@ -3159,6 +3199,8 @@ components:
         sourceCatalogId:
           type: string
           format: uuid
+        geography:
+          $ref: "#/components/schemas/Geography"
     WebBackendConnectionCreate:
       type: object
       required:
@@ -3206,6 +3248,8 @@ components:
         sourceCatalogId:
           type: string
           format: uuid
+        geography:
+          $ref: "#/components/schemas/Geography"
     ConnectionStateCreateOrUpdate:
       type: object
       required:
@@ -3256,6 +3300,8 @@ components:
         sourceCatalogId:
           type: string
           format: uuid
+        geography:
+          $ref: "#/components/schemas/Geography"
     WebBackendConnectionUpdate:
       type: object
       description: Used to apply a patch-style update to a connection, which means that null properties remain unchanged
@@ -3298,6 +3344,8 @@ components:
         sourceCatalogId:
           type: string
           format: uuid
+        geography:
+          $ref: "#/components/schemas/Geography"
     ConnectionRead:
       type: object
       required:
@@ -3345,6 +3393,8 @@ components:
         sourceCatalogId:
           type: string
           format: uuid
+        geography:
+          $ref: "#/components/schemas/Geography"
     ConnectionSearch:
       type: object
       properties:
@@ -4619,6 +4669,8 @@ components:
           format: uuid
         catalogDiff:
           $ref: "#/components/schemas/CatalogDiff"
+        geography:
+          $ref: "#/components/schemas/Geography"
     WebBackendConnectionReadList:
       type: object
       required:

airbyte-bootloader/src/main/java/io/airbyte/bootloader/BootloaderApp.java

@@ -12,6 +12,7 @@
 import io.airbyte.commons.version.Version;
 import io.airbyte.config.Configs;
 import io.airbyte.config.EnvConfigs;
+import io.airbyte.config.Geography;
 import io.airbyte.config.StandardWorkspace;
 import io.airbyte.config.init.ApplyDefinitionsHelper;
 import io.airbyte.config.init.DefinitionsProvider;
@@ -286,7 +287,8 @@ private static void createWorkspaceIfNoneExists(final ConfigRepository configRep
         .withSlug(workspaceId.toString())
         .withInitialSetupComplete(false)
         .withDisplaySetupWizard(true)
-        .withTombstone(false);
+        .withTombstone(false)
+        .withDefaultGeography(Geography.AUTO);
     configRepository.writeStandardWorkspace(workspace);
   }
 

airbyte-bootloader/src/test/java/io/airbyte/bootloader/BootloaderAppTest.java

@@ -21,6 +21,7 @@
 import io.airbyte.commons.version.AirbyteVersion;
 import io.airbyte.commons.version.Version;
 import io.airbyte.config.Configs;
+import io.airbyte.config.Geography;
 import io.airbyte.config.SourceConnection;
 import io.airbyte.config.StandardWorkspace;
 import io.airbyte.config.init.YamlSeedConfigPersistence;
@@ -222,7 +223,8 @@ void testBootloaderAppRunSecretMigration() throws Exception {
           .withSlug("wSlug")
           .withEmail("email@mail.com")
           .withTombstone(false)
-          .withInitialSetupComplete(false));
+          .withInitialSetupComplete(false)
+          .withDefaultGeography(Geography.AUTO));
       final UUID sourceId = UUID.randomUUID();
       configRepository.writeSourceConnectionNoSecrets(new SourceConnection()
           .withSourceDefinitionId(UUID.fromString("e7778cfc-e97c-4458-9ecb-b4f2bba8946c")) // Facebook Marketing

airbyte-config/config-models/src/main/resources/types/Geography.yaml

@@ -0,0 +1,10 @@
+---
+"$schema": http://json-schema.org/draft-07/schema#
+"$id": https://github.com/airbytehq/airbyte/blob/master/airbyte-config/models/src/main/resources/types/Geography.yaml
+title: Geography
+description: Geography Setting
+type: string
+enum:
+  - auto
+  - us
+  - eu

airbyte-config/config-models/src/main/resources/types/StandardSync.yaml

@@ -11,6 +11,7 @@ required:
   - catalog
   - manual
   - namespaceDefinition
+  - geography
 additionalProperties: false
 properties:
   namespaceDefinition:
@@ -114,3 +115,5 @@ properties:
     format: uuid
   resourceRequirements:
     "$ref": ResourceRequirements.yaml
+  geography:
+    "$ref": Geography.yaml

airbyte-config/config-models/src/main/resources/types/StandardWorkspace.yaml

@@ -9,6 +9,7 @@ required:
   - name
   - slug
   - initialSetupComplete
+  - defaultGeography
 additionalProperties: false
 properties:
   workspaceId:
@@ -47,3 +48,5 @@ properties:
     type: boolean
   feedbackDone:
     type: boolean
+  defaultGeography:
+    "$ref": Geography.yaml

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

@@ -1087,12 +1087,16 @@ private void writeStandardSync(final List<StandardSync> configs, final DSLContex
             .set(CONNECTION.MANUAL, standardSync.getManual())
             .set(CONNECTION.SCHEDULE_TYPE,
                 standardSync.getScheduleType() == null ? null
-                    : Enums.toEnum(standardSync.getScheduleType().value(), io.airbyte.db.instance.configs.jooq.generated.enums.ScheduleType.class)
+                    : Enums.toEnum(standardSync.getScheduleType().value(),
+                        io.airbyte.db.instance.configs.jooq.generated.enums.ScheduleType.class)
                         .orElseThrow())
             .set(CONNECTION.SCHEDULE_DATA, JSONB.valueOf(Jsons.serialize(standardSync.getScheduleData())))
-            .set(CONNECTION.RESOURCE_REQUIREMENTS, JSONB.valueOf(Jsons.serialize(standardSync.getResourceRequirements())))
+            .set(CONNECTION.RESOURCE_REQUIREMENTS,
+                JSONB.valueOf(Jsons.serialize(standardSync.getResourceRequirements())))
             .set(CONNECTION.UPDATED_AT, timestamp)
             .set(CONNECTION.SOURCE_CATALOG_ID, standardSync.getSourceCatalogId())
+            .set(CONNECTION.GEOGRAPHY, Enums.toEnum(standardSync.getGeography().value(),
+                io.airbyte.db.instance.configs.jooq.generated.enums.GeographyType.class).orElseThrow())
             .where(CONNECTION.ID.eq(standardSync.getConnectionId()))
             .execute();
 
@@ -1126,11 +1130,15 @@ private void writeStandardSync(final List<StandardSync> configs, final DSLContex
             .set(CONNECTION.MANUAL, standardSync.getManual())
             .set(CONNECTION.SCHEDULE_TYPE,
                 standardSync.getScheduleType() == null ? null
-                    : Enums.toEnum(standardSync.getScheduleType().value(), io.airbyte.db.instance.configs.jooq.generated.enums.ScheduleType.class)
+                    : Enums.toEnum(standardSync.getScheduleType().value(),
+                        io.airbyte.db.instance.configs.jooq.generated.enums.ScheduleType.class)
                         .orElseThrow())
             .set(CONNECTION.SCHEDULE_DATA, JSONB.valueOf(Jsons.serialize(standardSync.getScheduleData())))
-            .set(CONNECTION.RESOURCE_REQUIREMENTS, JSONB.valueOf(Jsons.serialize(standardSync.getResourceRequirements())))
+            .set(CONNECTION.RESOURCE_REQUIREMENTS,
+                JSONB.valueOf(Jsons.serialize(standardSync.getResourceRequirements())))
             .set(CONNECTION.SOURCE_CATALOG_ID, standardSync.getSourceCatalogId())
+            .set(CONNECTION.GEOGRAPHY, Enums.toEnum(standardSync.getGeography().value(),
+                io.airbyte.db.instance.configs.jooq.generated.enums.GeographyType.class).orElseThrow())
             .set(CONNECTION.CREATED_AT, timestamp)
             .set(CONNECTION.UPDATED_AT, timestamp)
             .execute();

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

@@ -18,6 +18,7 @@
 import io.airbyte.config.ActorDefinitionResourceRequirements;
 import io.airbyte.config.DestinationConnection;
 import io.airbyte.config.DestinationOAuthParameter;
+import io.airbyte.config.Geography;
 import io.airbyte.config.JobSyncConfig.NamespaceDefinitionType;
 import io.airbyte.config.Notification;
 import io.airbyte.config.ResourceRequirements;
@@ -60,16 +61,20 @@ public static StandardSync buildStandardSync(final Record record, final List<UUI
         .withCatalog(
             Jsons.deserialize(record.get(CONNECTION.CATALOG).data(), ConfiguredAirbyteCatalog.class))
         .withStatus(
-            record.get(CONNECTION.STATUS) == null ? null : Enums.toEnum(record.get(CONNECTION.STATUS, String.class), Status.class).orElseThrow())
+            record.get(CONNECTION.STATUS) == null ? null
+                : Enums.toEnum(record.get(CONNECTION.STATUS, String.class), Status.class).orElseThrow())
         .withSchedule(Jsons.deserialize(record.get(CONNECTION.SCHEDULE).data(), Schedule.class))
         .withManual(record.get(CONNECTION.MANUAL))
         .withScheduleType(record.get(CONNECTION.SCHEDULE_TYPE) == null ? null
             : Enums.toEnum(record.get(CONNECTION.SCHEDULE_TYPE, String.class), ScheduleType.class).orElseThrow())
         .withScheduleData(
-            record.get(CONNECTION.SCHEDULE_DATA) == null ? null : Jsons.deserialize(record.get(CONNECTION.SCHEDULE_DATA).data(), ScheduleData.class))
+            record.get(CONNECTION.SCHEDULE_DATA) == null ? null
+                : Jsons.deserialize(record.get(CONNECTION.SCHEDULE_DATA).data(), ScheduleData.class))
         .withOperationIds(connectionOperationId)
-        .withResourceRequirements(Jsons.deserialize(record.get(CONNECTION.RESOURCE_REQUIREMENTS).data(), ResourceRequirements.class))
-        .withSourceCatalogId(record.get(CONNECTION.SOURCE_CATALOG_ID));
+        .withResourceRequirements(
+            Jsons.deserialize(record.get(CONNECTION.RESOURCE_REQUIREMENTS).data(), ResourceRequirements.class))
+        .withSourceCatalogId(record.get(CONNECTION.SOURCE_CATALOG_ID))
+        .withGeography(Enums.toEnum(record.get(CONNECTION.GEOGRAPHY, String.class), Geography.class).orElseThrow());
   }
 
   public static StandardWorkspace buildStandardWorkspace(final Record record) {
@@ -92,7 +97,9 @@ public static StandardWorkspace buildStandardWorkspace(final Record record) {
         .withTombstone(record.get(WORKSPACE.TOMBSTONE))
         .withNotifications(notificationList)
         .withFirstCompletedSync(record.get(WORKSPACE.FIRST_SYNC_COMPLETE))
-        .withFeedbackDone(record.get(WORKSPACE.FEEDBACK_COMPLETE));
+        .withFeedbackDone(record.get(WORKSPACE.FEEDBACK_COMPLETE))
+        .withDefaultGeography(
+            Enums.toEnum(record.get(WORKSPACE.GEOGRAPHY, String.class), Geography.class).orElseThrow());
   }
 
   public static SourceConnection buildSourceConnection(final Record record) {

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

@@ -12,6 +12,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import io.airbyte.config.ConfigSchema;
 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;
@@ -146,7 +147,8 @@ void writeSourceWithSourceConnection(final ConfigPersistence configPersistence,
         .withWorkspaceId(workspaceId)
         .withName(CANNOT_BE_NULL)
         .withSlug(CANNOT_BE_NULL)
-        .withInitialSetupComplete(true);
+        .withInitialSetupComplete(true)
+        .withDefaultGeography(Geography.AUTO);
     configPersistence.writeConfig(ConfigSchema.STANDARD_WORKSPACE, workspaceId.toString(), workspace);
 
     final SourceConnection sourceConnection = new SourceConnection()
@@ -172,7 +174,8 @@ void writeDestinationWithDestinationConnection(final ConfigPersistence configPer
         .withWorkspaceId(workspaceId)
         .withName(CANNOT_BE_NULL)
         .withSlug(CANNOT_BE_NULL)
-        .withInitialSetupComplete(true);
+        .withInitialSetupComplete(true)
+        .withDefaultGeography(Geography.AUTO);
     configPersistence.writeConfig(ConfigSchema.STANDARD_WORKSPACE, workspaceId.toString(), workspace);
 
     final DestinationConnection destinationConnection = new DestinationConnection()

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

@@ -17,6 +17,7 @@
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.config.ConfigSchema;
 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;
@@ -119,7 +120,8 @@ void testNoUpdateForUsedConnector() throws Exception {
         .withWorkspaceId(s3Connection.getWorkspaceId())
         .withName("workspace")
         .withSlug("slug")
-        .withInitialSetupComplete(true);
+        .withInitialSetupComplete(true)
+        .withDefaultGeography(Geography.AUTO);
     configPersistence.writeConfig(ConfigSchema.STANDARD_WORKSPACE, standardWorkspace.getWorkspaceId().toString(), standardWorkspace);
     configPersistence.writeConfig(ConfigSchema.DESTINATION_CONNECTION, s3Connection.getDestinationId().toString(), s3Connection);
     final SourceConnection githubConnection = new SourceConnection()