Add deployment mode to segment event tracking (#8578)

duffelhq · Aug 31, 2023 · a4e4b65 · a4e4b65
1 parent f0d19b8
commit a4e4b65
Show file tree

Hide file tree

Showing 11 changed files with 46 additions and 20 deletions.
diff --git a/airbyte-analytics/src/main/java/io/airbyte/analytics/SegmentTrackingClient.java b/airbyte-analytics/src/main/java/io/airbyte/analytics/SegmentTrackingClient.java
@@ -57,6 +57,7 @@ public class SegmentTrackingClient implements TrackingClient {
   private static final String AIRBYTE_TRACKED_AT = "tracked_at";
   protected static final String UNKNOWN = "unknown";
   protected static final String AIRBYTE_DEPLOYMENT_ID = "deployment_id";
+  protected static final String AIRBYTE_DEPLOYMENT_MODE = "deployment_mode";
 
   // Analytics is threadsafe.
   private final Analytics analytics;
@@ -138,6 +139,7 @@ public void track(@Nullable final UUID workspaceId, final String action, final M
     mapCopy.put(AIRBYTE_VERSION_KEY, trackingIdentity.getAirbyteVersion().serialize());
     mapCopy.put(CUSTOMER_ID_KEY, trackingIdentity.getCustomerId());
     mapCopy.put(AIRBYTE_DEPLOYMENT_ID, deployment.getDeploymentId());
+    mapCopy.put(AIRBYTE_DEPLOYMENT_MODE, deployment.getDeploymentMode());
     mapCopy.put(AIRBYTE_TRACKED_AT, Instant.now().toString());
     if (!metadata.isEmpty()) {
       trackingIdentity.getEmail().ifPresent(email -> mapCopy.put("email", email));

diff --git a/airbyte-analytics/src/test/java/io/airbyte/analytics/SegmentTrackingClientTest.java b/airbyte-analytics/src/test/java/io/airbyte/analytics/SegmentTrackingClientTest.java
@@ -6,6 +6,7 @@
 
 import static io.airbyte.analytics.SegmentTrackingClient.AIRBYTE_ANALYTIC_SOURCE_HEADER;
 import static io.airbyte.analytics.SegmentTrackingClient.AIRBYTE_DEPLOYMENT_ID;
+import static io.airbyte.analytics.SegmentTrackingClient.AIRBYTE_DEPLOYMENT_MODE;
 import static io.airbyte.analytics.SegmentTrackingClient.AIRBYTE_SOURCE;
 import static io.airbyte.analytics.SegmentTrackingClient.AIRBYTE_VERSION_KEY;
 import static io.airbyte.analytics.SegmentTrackingClient.UNKNOWN;
@@ -119,7 +120,8 @@ void testTrack() {
         ImmutableMap.of(AIRBYTE_VERSION_KEY, AIRBYTE_VERSION.serialize(),
             "user_id", IDENTITY.getCustomerId(),
             AIRBYTE_SOURCE, UNKNOWN,
-            AIRBYTE_DEPLOYMENT_ID, DEPLOYMENT.getDeploymentId());
+            AIRBYTE_DEPLOYMENT_ID, DEPLOYMENT.getDeploymentId(),
+            AIRBYTE_DEPLOYMENT_MODE, DEPLOYMENT.getDeploymentMode());
 
     segmentTrackingClient.track(WORKSPACE_ID, JUMP);
 
@@ -139,7 +141,8 @@ void testTrackWithMetadata() {
         "height", "80 meters",
         "user_id", IDENTITY.getCustomerId(),
         AIRBYTE_SOURCE, UNKNOWN,
-        AIRBYTE_DEPLOYMENT_ID, DEPLOYMENT.getDeploymentId());
+        AIRBYTE_DEPLOYMENT_ID, DEPLOYMENT.getDeploymentId(),
+        AIRBYTE_DEPLOYMENT_MODE, DEPLOYMENT.getDeploymentMode());
 
     segmentTrackingClient.track(WORKSPACE_ID, JUMP, metadata);
 

diff --git a/.../api/server/helpers/ApiTrackingHelpers.kt → .../server/apiTracking/ApiTrackingHelpers.kt b/.../api/server/helpers/ApiTrackingHelpers.kt → .../server/apiTracking/ApiTrackingHelpers.kt
@@ -1,7 +1,7 @@
 /*
  * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  */
-package io.airbyte.api.server.helpers
+package io.airbyte.api.server.apiTracking
 
 import io.airbyte.analytics.Deployment
 import io.airbyte.analytics.TrackingClientSingleton
@@ -31,23 +31,29 @@ private val USER_ID = "user_id"
 private val ENDPOINT = "endpoint"
 private val OPERATION = "operation"
 private val STATUS_CODE = "status_code"
-private val API_VERSION = "api_version"
 private val WORKSPACE = "workspace"
 
 fun setupTrackingClient(
-  airbyteVersion: AirbyteVersion?,
-  trackingStrategy: Configs.TrackingStrategy?,
+  airbyteVersion: AirbyteVersion,
+  deploymentMode: Configs.DeploymentMode,
+  trackingStrategy: Configs.TrackingStrategy,
+  workerEnvironment: Configs.WorkerEnvironment,
 ) {
+  log.debug("deployment mode: $deploymentMode")
+  log.debug("airbyte version: $airbyteVersion")
+  log.debug("tracking strategy: $trackingStrategy")
+  log.debug("worker environment: $workerEnvironment")
+
   // fake a deployment UUID until we want to have the public api server talking to the database
   // directly
   val deploymentId = UUID.randomUUID()
   log.info("setting up tracking client with deploymentId: $deploymentId")
   TrackingClientSingleton.initializeWithoutDatabase(
     trackingStrategy,
     Deployment(
-      Configs.DeploymentMode.OSS,
+      deploymentMode,
       deploymentId,
-      Configs.WorkerEnvironment.KUBERNETES,
+      workerEnvironment,
     ),
     airbyteVersion,
   )
@@ -58,15 +64,13 @@ fun track(
   endpointPath: String?,
   httpOperation: String?,
   httpStatusCode: Int,
-  apiVersion: String?,
   workspaceId: Optional<UUID>,
 ) {
   val payload = mutableMapOf(
     Pair(USER_ID, userId),
     Pair(ENDPOINT, endpointPath),
     Pair(OPERATION, httpOperation),
     Pair(STATUS_CODE, httpStatusCode),
-    Pair(API_VERSION, apiVersion),
   )
   if (workspaceId.isPresent) {
     payload[WORKSPACE] = workspaceId.get().toString()

diff --git a/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/apiTracking/TrackingClient.kt b/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/apiTracking/TrackingClient.kt
@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+ */
+package io.airbyte.api.server.apiTracking
+
+import io.airbyte.config.Configs
+import io.airbyte.config.EnvConfigs
+import jakarta.inject.Singleton
+
+/**
+ * Used to track usage of the airbyte-api-server.
+ */
+@Singleton
+class TrackingClient {
+  val configs: Configs = EnvConfigs()
+
+  init {
+    setupTrackingClient(configs.airbyteVersion, configs.deploymentMode, configs.trackingStrategy, configs.workerEnvironment)
+  }
+}
diff --git a/...byte/api/server/helpers/TrackingHelper.kt → .../api/server/apiTracking/TrackingHelper.kt b/...byte/api/server/helpers/TrackingHelper.kt → .../api/server/apiTracking/TrackingHelper.kt
@@ -2,7 +2,7 @@
  * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  */
 
-package io.airbyte.api.server.helpers
+package io.airbyte.api.server.apiTracking
 
 import io.micronaut.http.HttpStatus
 import org.zalando.problem.AbstractThrowableProblem
@@ -16,15 +16,13 @@ import javax.ws.rs.core.Response
  * Todo: This should be a middleware through a micronaut annotation so that we do not need to add this wrapper functions around all our calls.
  */
 object TrackingHelper {
-  private val API_VERSION = System.getenv("AIRBYTE_VERSION") ?: "unknown"
   private fun trackSuccess(endpointPath: String, httpOperation: String, userId: UUID, workspaceId: Optional<UUID>) {
     val statusCode = Response.Status.OK.statusCode
     track(
       userId,
       endpointPath,
       httpOperation,
       statusCode,
-      API_VERSION,
       workspaceId,
     )
   }
@@ -62,7 +60,6 @@ object TrackingHelper {
         endpointPath,
         httpOperation,
         statusCode,
-        API_VERSION,
         Optional.empty(),
       )
     }

diff --git a/...yte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/ConnectionsController.kt b/...yte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/ConnectionsController.kt
@@ -13,14 +13,14 @@ import io.airbyte.api.client.model.generated.AirbyteCatalog
 import io.airbyte.api.client.model.generated.AirbyteStreamAndConfiguration
 import io.airbyte.api.client.model.generated.DestinationSyncMode
 import io.airbyte.api.client.model.generated.SourceDiscoverSchemaRead
+import io.airbyte.api.server.apiTracking.TrackingHelper
 import io.airbyte.api.server.constants.CONNECTIONS_PATH
 import io.airbyte.api.server.constants.CONNECTIONS_WITH_ID_PATH
 import io.airbyte.api.server.constants.DELETE
 import io.airbyte.api.server.constants.GET
 import io.airbyte.api.server.constants.POST
 import io.airbyte.api.server.constants.PUT
 import io.airbyte.api.server.helpers.AirbyteCatalogHelper
-import io.airbyte.api.server.helpers.TrackingHelper
 import io.airbyte.api.server.helpers.getLocalUserInfoIfNull
 import io.airbyte.api.server.services.ConnectionService
 import io.airbyte.api.server.services.DestinationService

diff --git a/...te-api-server/src/main/kotlin/io/airbyte/api/server/controllers/DestinationsController.kt b/...te-api-server/src/main/kotlin/io/airbyte/api/server/controllers/DestinationsController.kt
@@ -9,6 +9,7 @@ import io.airbyte.airbyte_api.generated.DestinationsApi
 import io.airbyte.airbyte_api.model.generated.DestinationCreateRequest
 import io.airbyte.airbyte_api.model.generated.DestinationPatchRequest
 import io.airbyte.airbyte_api.model.generated.DestinationPutRequest
+import io.airbyte.api.server.apiTracking.TrackingHelper
 import io.airbyte.api.server.constants.DELETE
 import io.airbyte.api.server.constants.DESTINATIONS_PATH
 import io.airbyte.api.server.constants.DESTINATIONS_WITH_ID_PATH
@@ -17,7 +18,6 @@ import io.airbyte.api.server.constants.GET
 import io.airbyte.api.server.constants.PATCH
 import io.airbyte.api.server.constants.POST
 import io.airbyte.api.server.constants.PUT
-import io.airbyte.api.server.helpers.TrackingHelper
 import io.airbyte.api.server.helpers.getIdFromName
 import io.airbyte.api.server.helpers.getLocalUserInfoIfNull
 import io.airbyte.api.server.helpers.removeDestinationType

diff --git a/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/JobsController.kt b/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/JobsController.kt
@@ -11,13 +11,13 @@ import io.airbyte.airbyte_api.model.generated.JobStatusEnum
 import io.airbyte.airbyte_api.model.generated.JobTypeEnum
 import io.airbyte.api.client.model.generated.JobListForWorkspacesRequestBody.OrderByFieldEnum
 import io.airbyte.api.client.model.generated.JobListForWorkspacesRequestBody.OrderByMethodEnum
+import io.airbyte.api.server.apiTracking.TrackingHelper
 import io.airbyte.api.server.constants.DELETE
 import io.airbyte.api.server.constants.GET
 import io.airbyte.api.server.constants.JOBS_PATH
 import io.airbyte.api.server.constants.JOBS_WITH_ID_PATH
 import io.airbyte.api.server.constants.POST
 import io.airbyte.api.server.filters.JobsFilter
-import io.airbyte.api.server.helpers.TrackingHelper
 import io.airbyte.api.server.helpers.getLocalUserInfoIfNull
 import io.airbyte.api.server.problems.BadRequestProblem
 import io.airbyte.api.server.problems.UnprocessableEntityProblem

diff --git a/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/SourcesController.kt b/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/SourcesController.kt
@@ -10,6 +10,7 @@ import io.airbyte.airbyte_api.model.generated.InitiateOauthRequest
 import io.airbyte.airbyte_api.model.generated.SourceCreateRequest
 import io.airbyte.airbyte_api.model.generated.SourcePatchRequest
 import io.airbyte.airbyte_api.model.generated.SourcePutRequest
+import io.airbyte.api.server.apiTracking.TrackingHelper
 import io.airbyte.api.server.constants.DELETE
 import io.airbyte.api.server.constants.GET
 import io.airbyte.api.server.constants.PATCH
@@ -18,7 +19,6 @@ import io.airbyte.api.server.constants.PUT
 import io.airbyte.api.server.constants.SOURCES_PATH
 import io.airbyte.api.server.constants.SOURCES_WITH_ID_PATH
 import io.airbyte.api.server.constants.SOURCE_TYPE
-import io.airbyte.api.server.helpers.TrackingHelper
 import io.airbyte.api.server.helpers.getIdFromName
 import io.airbyte.api.server.helpers.getLocalUserInfoIfNull
 import io.airbyte.api.server.helpers.removeSourceTypeNode

diff --git a/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/StreamsController.kt b/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/StreamsController.kt
@@ -13,9 +13,9 @@ import io.airbyte.airbyte_api.model.generated.StreamProperties
 import io.airbyte.api.client.model.generated.AirbyteStreamAndConfiguration
 import io.airbyte.api.client.model.generated.DestinationSyncMode
 import io.airbyte.api.client.model.generated.SyncMode
+import io.airbyte.api.server.apiTracking.TrackingHelper
 import io.airbyte.api.server.constants.GET
 import io.airbyte.api.server.constants.STREAMS_PATH
-import io.airbyte.api.server.helpers.TrackingHelper
 import io.airbyte.api.server.helpers.getLocalUserInfoIfNull
 import io.airbyte.api.server.problems.UnexpectedProblem
 import io.airbyte.api.server.services.DestinationService

diff --git a/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/WorkspacesController.kt b/airbyte-api-server/src/main/kotlin/io/airbyte/api/server/controllers/WorkspacesController.kt
@@ -9,12 +9,12 @@ import io.airbyte.airbyte_api.model.generated.WorkspaceCreateRequest
 import io.airbyte.airbyte_api.model.generated.WorkspaceOAuthCredentialsRequest
 import io.airbyte.airbyte_api.model.generated.WorkspaceResponse
 import io.airbyte.airbyte_api.model.generated.WorkspaceUpdateRequest
+import io.airbyte.api.server.apiTracking.TrackingHelper
 import io.airbyte.api.server.constants.DELETE
 import io.airbyte.api.server.constants.GET
 import io.airbyte.api.server.constants.POST
 import io.airbyte.api.server.constants.WORKSPACES_PATH
 import io.airbyte.api.server.constants.WORKSPACES_WITH_ID_PATH
-import io.airbyte.api.server.helpers.TrackingHelper
 import io.airbyte.api.server.helpers.getLocalUserInfoIfNull
 import io.airbyte.api.server.services.UserService
 import io.airbyte.api.server.services.WorkspaceService