track by workst pace id

airbytehq · Jul 19, 2021 · 7062113 · 7062113
1 parent 5af32fc
commit 7062113
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 34 deletions.
diff --git a/airbyte-analytics/src/main/java/io/airbyte/analytics/LoggingTrackingClient.java b/airbyte-analytics/src/main/java/io/airbyte/analytics/LoggingTrackingClient.java
@@ -26,6 +26,7 @@
 
 import java.util.Collections;
 import java.util.Map;
+import java.util.UUID;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -41,22 +42,22 @@ public LoggingTrackingClient(Supplier<TrackingIdentity> identitySupplier) {
   }
 
   @Override
-  public void identify() {
+  public void identify(UUID workspaceId) {
     LOGGER.info("identify. userId: {}", identitySupplier.get().getCustomerId());
   }
 
   @Override
-  public void alias(String previousCustomerId) {
+  public void alias(UUID workspaceId, String previousCustomerId) {
     LOGGER.info("merge. userId: {} previousUserId: {}", identitySupplier.get().getCustomerId(), previousCustomerId);
   }
 
   @Override
-  public void track(String action) {
-    track(action, Collections.emptyMap());
+  public void track(UUID workspaceId, String action) {
+    track(workspaceId, action, Collections.emptyMap());
   }
 
   @Override
-  public void track(String action, Map<String, Object> metadata) {
+  public void track(UUID workspaceId, String action, Map<String, Object> metadata) {
     LOGGER.info("track. version: {}, userId: {}, action: {}, metadata: {}",
         identitySupplier.get().getAirbyteVersion(),
         identitySupplier.get().getCustomerId(),

diff --git a/airbyte-analytics/src/main/java/io/airbyte/analytics/SegmentTrackingClient.java b/airbyte-analytics/src/main/java/io/airbyte/analytics/SegmentTrackingClient.java
@@ -35,6 +35,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.UUID;
 import java.util.function.Supplier;
 
 public class SegmentTrackingClient implements TrackingClient {
@@ -67,7 +68,7 @@ public SegmentTrackingClient(final Supplier<TrackingIdentity> identitySupplier,
   }
 
   @Override
-  public void identify() {
+  public void identify(UUID workspaceId) {
     final TrackingIdentity trackingIdentity = identitySupplier.get();
     final Map<String, Object> identityMetadata = new HashMap<>();
 
@@ -93,17 +94,17 @@ public void identify() {
   }
 
   @Override
-  public void alias(String previousCustomerId) {
+  public void alias(UUID workspaceId, String previousCustomerId) {
     analytics.enqueue(AliasMessage.builder(previousCustomerId).userId(identitySupplier.get().getCustomerId().toString()));
   }
 
   @Override
-  public void track(String action) {
-    track(action, Collections.emptyMap());
+  public void track(UUID workspaceId, String action) {
+    track(workspaceId, action, Collections.emptyMap());
   }
 
   @Override
-  public void track(String action, Map<String, Object> metadata) {
+  public void track(UUID workspaceId, String action, Map<String, Object> metadata) {
     final Map<String, Object> mapCopy = new HashMap<>(metadata);
     final TrackingIdentity trackingIdentity = identitySupplier.get();
     mapCopy.put(AIRBYTE_VERSION_KEY, trackingIdentity.getAirbyteVersion());

diff --git a/airbyte-analytics/src/main/java/io/airbyte/analytics/TrackingClient.java b/airbyte-analytics/src/main/java/io/airbyte/analytics/TrackingClient.java
@@ -25,15 +25,16 @@
 package io.airbyte.analytics;
 
 import java.util.Map;
+import java.util.UUID;
 
 public interface TrackingClient {
 
-  void identify();
+  void identify(UUID workspaceId);
 
-  void alias(String previousCustomerId);
+  void alias(UUID workspaceId, String previousCustomerId);
 
-  void track(String action);
+  void track(UUID workspaceId, String action);
 
-  void track(String action, Map<String, Object> metadata);
+  void track(UUID workspaceId, String action, Map<String, Object> metadata);
 
 }
diff --git a/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java b/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java
@@ -33,11 +33,9 @@
 import io.airbyte.config.EnvConfigs;
 import io.airbyte.config.StandardWorkspace;
 import io.airbyte.config.helpers.LogClientSingleton;
-import io.airbyte.config.persistence.ConfigNotFoundException;
 import io.airbyte.config.persistence.ConfigPersistence;
 import io.airbyte.config.persistence.ConfigPersistenceBuilder;
 import io.airbyte.config.persistence.ConfigRepository;
-import io.airbyte.config.persistence.PersistenceConstants;
 import io.airbyte.db.Database;
 import io.airbyte.db.Databases;
 import io.airbyte.scheduler.client.DefaultSchedulerJobClient;
@@ -110,8 +108,6 @@ public ServerApp(final ConfigRepository configRepository,
   }
 
   public void start() throws Exception {
-    TrackingClientSingleton.get().identify();
-
     Server server = new Server(PORT);
 
     ServletContextHandler handler = new ServletContextHandler();
@@ -190,28 +186,37 @@ private static void createDeploymentIfNoneExists(final JobPersistence jobPersist
 
   private static void setCustomerIdIfNotSet(final ConfigRepository configRepository) throws InterruptedException {
     StandardWorkspace workspace = null;
+  }
+
+  private static void createWorkspaceIfNoneExists(final ConfigRepository configRepository)
+      throws JsonValidationException, IOException, InterruptedException {
+    if (!configRepository.listStandardWorkspaces(true).isEmpty()) {
+      LOGGER.info("workspace already exists for the deployment.");
+      return;
+    }
 
+    final UUID workspaceId = UUID.randomUUID();
     // retry until the workspace is available / waits for file config initialization
-    while (workspace == null) {
+    while (true) {
       try {
-        workspace = configRepository.getStandardWorkspace(PersistenceConstants.DEFAULT_WORKSPACE_ID, true);
-
-        if (workspace.getCustomerId() == null) {
-          final UUID customerId = UUID.randomUUID();
-          LOGGER.info("customerId not set for workspace. Setting it to " + customerId);
-          workspace.setCustomerId(customerId);
-
-          configRepository.writeStandardWorkspace(workspace);
-        } else {
-          LOGGER.info("customerId already set for workspace: " + workspace.getCustomerId());
-        }
-      } catch (ConfigNotFoundException e) {
-        LOGGER.error("Could not find workspace with id: " + PersistenceConstants.DEFAULT_WORKSPACE_ID, e);
+        final StandardWorkspace workspace = new StandardWorkspace()
+            .withWorkspaceId(workspaceId)
+            .withCustomerId(UUID.randomUUID())
+            .withName(workspaceId.toString())
+            .withSlug(workspaceId.toString())
+            .withInitialSetupComplete(false)
+            .withDisplaySetupWizard(true)
+            .withTombstone(false);
+        configRepository.writeStandardWorkspace(workspace);
+        break;
+      } catch (IOException e) {
+        LOGGER.warn("Failed to create a starter workspace", e);
         Thread.sleep(1000);
-      } catch (JsonValidationException | IOException e) {
+      } catch (JsonValidationException e) {
         throw new RuntimeException(e);
       }
     }
+    TrackingClientSingleton.get().identify(workspaceId);
   }
 
   public static void runServer(final Set<ContainerRequestFilter> requestFilters,
@@ -227,7 +232,7 @@ public static void runServer(final Set<ContainerRequestFilter> requestFilters,
 
     // hack: upon installation we need to assign a random customerId so that when
     // tracking we can associate all action with the correct anonymous id.
-    setCustomerIdIfNotSet(configRepository);
+    createWorkspaceIfNoneExists(configRepository);
 
     TrackingClientSingleton.initialize(
         configs.getTrackingStrategy(),