router service to route sync to expanded dataplane (#7476)

airbytehq · Jun 23, 2023 · f3b717b · f3b717b
1 parent fe230ce
commit f3b717b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/...-commons-temporal/src/main/java/io/airbyte/commons/temporal/scheduling/RouterService.java b/...-commons-temporal/src/main/java/io/airbyte/commons/temporal/scheduling/RouterService.java
@@ -48,7 +48,11 @@ public String getTaskQueue(final UUID connectionId, final TemporalJobType jobTyp
     final Geography geography = configRepository.getGeographyForConnection(connectionId);
     final UUID workspaceId = configRepository.getStandardWorkspaceFromConnection(connectionId, false).getWorkspaceId();
     if (featureFlagClient.boolVariation(ShouldRunOnGkeDataplane.INSTANCE, new Workspace(workspaceId))) {
-      return taskQueueMapper.getTaskQueueFlagged(geography, jobType);
+      if (featureFlagClient.boolVariation(ShouldRunOnExpandedGkeDataplane.INSTANCE, new Workspace(workspaceId))) {
+        return taskQueueMapper.getTaskQueueExpanded(geography, jobType);
+      } else {
+        return taskQueueMapper.getTaskQueueFlagged(geography, jobType);
+      }
     } else {
       return taskQueueMapper.getTaskQueue(geography, jobType);
     }

diff --git a/airbyte-workers/src/test/java/io/airbyte/commons/temporal/scheduling/RouterServiceTest.java b/airbyte-workers/src/test/java/io/airbyte/commons/temporal/scheduling/RouterServiceTest.java
@@ -105,6 +105,21 @@ void testGetTaskQueueBehindFlag() throws IOException, ConfigNotFoundException {
     assertEquals(EU_FLAGGED_TASK_QUEUE, routerService.getTaskQueue(CONNECTION_ID, TemporalJobType.SYNC));
   }
 
+  @Test
+  void testGetTaskQueueBehindExpandedFlag() throws IOException, ConfigNotFoundException {
+    Mockito.when(mockFeatureFlagClient.boolVariation(ShouldRunOnGkeDataplane.INSTANCE, new Workspace(WORKSPACE_ID))).thenReturn(true);
+    Mockito.when(mockFeatureFlagClient.boolVariation(ShouldRunOnExpandedGkeDataplane.INSTANCE, new Workspace(WORKSPACE_ID))).thenReturn(true);
+
+    Mockito.when(mConfigRepository.getGeographyForConnection(CONNECTION_ID)).thenReturn(Geography.AUTO);
+    assertEquals(US_EXPANDED_TASK_QUEUE, routerService.getTaskQueue(CONNECTION_ID, TemporalJobType.SYNC));
+
+    Mockito.when(mConfigRepository.getGeographyForConnection(CONNECTION_ID)).thenReturn(Geography.US);
+    assertEquals(US_EXPANDED_TASK_QUEUE, routerService.getTaskQueue(CONNECTION_ID, TemporalJobType.SYNC));
+
+    Mockito.when(mConfigRepository.getGeographyForConnection(CONNECTION_ID)).thenReturn(Geography.EU);
+    assertEquals(EU_EXPANDED_TASK_QUEUE, routerService.getTaskQueue(CONNECTION_ID, TemporalJobType.SYNC));
+  }
+
   @Test
   void testGetWorkspaceTaskQueue() throws IOException, ConfigNotFoundException {
     Mockito.when(mockFeatureFlagClient.boolVariation(ShouldRunOnGkeDataplane.INSTANCE, new Workspace(WORKSPACE_ID))).thenReturn(false);