fix: add recent job filter for stream status query (#18515)

Author: subodh1810

airbyte-featureflag/src/main/kotlin/FlagDefinitions.kt

@@ -261,3 +261,5 @@ object BypassStiggEntitlementChecks : Permanent<Boolean>(key = "platform.bypass-
 object EnableDataWorkerUsage : Temporary<Boolean>(key = "platform.enable-data-worker-usage", default = false)
 
 object UseVerifiedDomainsForSsoActivate : Temporary<Boolean>(key = "platform.use-verified-domains-for-sso-activate", default = false)
+
+object UseOptimizedStreamStatusQuery : Temporary<Boolean>(key = "platform.use-optimized-stream-status-query", default = false)

airbyte-server/src/main/kotlin/io/airbyte/server/handlers/StreamStatusesHandler.kt

@@ -19,6 +19,9 @@ import io.airbyte.api.model.generated.StreamStatusUpdateRequestBody
 import io.airbyte.commons.server.handlers.JobHistoryHandler
 import io.airbyte.commons.server.handlers.helpers.StatsAggregationHelper
 import io.airbyte.config.Job
+import io.airbyte.featureflag.Connection
+import io.airbyte.featureflag.FeatureFlagClient
+import io.airbyte.featureflag.UseOptimizedStreamStatusQuery
 import io.airbyte.persistence.job.JobPersistence
 import io.airbyte.server.handlers.apidomainmapping.StreamStatusesMapper
 import io.airbyte.server.repositories.StreamStatusesRepository
@@ -34,6 +37,7 @@ open class StreamStatusesHandler(
   val mapper: StreamStatusesMapper,
   private val jobHistoryHandler: JobHistoryHandler,
   private val jobPersistence: JobPersistence,
+  private val featureFlagClient: FeatureFlagClient,
 ) {
   fun createStreamStatus(req: StreamStatusCreateRequestBody): StreamStatusRead {
     val model = mapper.map(req)
@@ -64,14 +68,25 @@ open class StreamStatusesHandler(
   }
 
   fun listStreamStatusPerRunState(req: ConnectionIdRequestBody): StreamStatusReadList {
-    val apiList =
-      repo
-        .findAllPerRunStateByConnectionId(req.connectionId)
-        .map { domain -> mapper.map(domain) }
+    val useOptimized = featureFlagClient.boolVariation(UseOptimizedStreamStatusQuery, Connection(req.connectionId))
 
+    val streamStatuses =
+      if (useOptimized) {
+        // Optimized: same query but with recency filter (last 100 jobs)
+        // Reduces query time from 2+ min to ~18 sec for high-volume connections
+        repo.findAllPerRunStateByConnectionIdWithRecentJobsFilter(req.connectionId, RECENT_JOBS_LIMIT)
+      } else {
+        repo.findAllPerRunStateByConnectionId(req.connectionId)
+      }
+
+    val apiList = streamStatuses.map { domain -> mapper.map(domain) }
     return StreamStatusReadList().streamStatuses(apiList)
   }
 
+  companion object {
+    const val RECENT_JOBS_LIMIT = 100
+  }
+
   fun mapStreamStatusToSyncReadResult(streamStatus: StreamStatusRead): ConnectionSyncResultRead {
     val jobStatus =
       if (streamStatus.runState == StreamStatusRunState.COMPLETE) {

airbyte-server/src/main/kotlin/io/airbyte/server/repositories/StreamStatusesRepository.kt

@@ -115,6 +115,41 @@ abstract class StreamStatusesRepository :
   )
   abstract fun findAllPerRunStateByConnectionId(connectionId: UUID): List<StreamStatus>
 
+  /**
+   * Returns the latest stream status per run state (and job type) for a connection,
+   * limited to recent jobs for better performance on high-volume connections.
+   *
+   * This is the same logic as findAllPerRunStateByConnectionId but with a recency filter
+   * that limits processing to the last N jobs, reducing query time from 2+ minutes to ~18 seconds
+   * for connections with millions of stream_statuses.
+   */
+  @Query(
+    """
+    SELECT DISTINCT ON (ss.stream_name, ss.stream_namespace, 1 /* computed run_state */, 2 /* computed incomplete_run_cause */, ss.job_type)
+      $STREAM_STATUS_WITH_FALLBACKS
+      FROM stream_statuses ss
+      JOIN jobs j on j.id = ss.job_id
+      WHERE ss.connection_id = :connectionId
+        AND ss.job_id >= (
+          SELECT COALESCE(MIN(id), 0)
+          FROM (
+            SELECT id FROM jobs
+            WHERE scope = CAST(:connectionId AS VARCHAR)
+              AND config_type IN ('sync', 'refresh')
+            ORDER BY id DESC
+            LIMIT :recentJobsLimit
+          ) recent_jobs
+        )
+      ORDER BY
+        ss.job_type, ss.stream_name, ss.stream_namespace,
+        1 /* computed run_state */,2 /* computed incomplete_run_cause */, ss.transitioned_at desc
+      """,
+  )
+  abstract fun findAllPerRunStateByConnectionIdWithRecentJobsFilter(
+    connectionId: UUID,
+    recentJobsLimit: Int,
+  ): List<StreamStatus>
+
   @Query(
     """
            SELECT

airbyte-server/src/test/kotlin/io/airbyte/server/handlers/StreamStatusesHandlerTest.kt

@@ -29,6 +29,9 @@ import io.airbyte.config.StreamSyncStats
 import io.airbyte.config.SyncStats
 import io.airbyte.db.instance.jobs.jooq.generated.enums.JobStreamStatusJobType
 import io.airbyte.db.instance.jobs.jooq.generated.enums.JobStreamStatusRunState
+import io.airbyte.featureflag.Connection
+import io.airbyte.featureflag.FeatureFlagClient
+import io.airbyte.featureflag.UseOptimizedStreamStatusQuery
 import io.airbyte.persistence.job.JobPersistence
 import io.airbyte.server.handlers.apidomainmapping.StreamStatusesMapper
 import io.airbyte.server.repositories.StreamStatusesRepository
@@ -51,14 +54,16 @@ internal class StreamStatusesHandlerTest {
   private lateinit var handler: StreamStatusesHandler
   private lateinit var jobPersistence: JobPersistence
   private lateinit var jobHistoryHandler: JobHistoryHandler
+  private lateinit var featureFlagClient: FeatureFlagClient
 
   @BeforeEach
   fun setup() {
     repo = mockk()
     mapper = mockk()
     jobHistoryHandler = mockk()
     jobPersistence = mockk()
-    handler = StreamStatusesHandler(repo, mapper, jobHistoryHandler, jobPersistence)
+    featureFlagClient = mockk()
+    handler = StreamStatusesHandler(repo, mapper, jobHistoryHandler, jobPersistence, featureFlagClient)
   }
 
   @Test
@@ -107,19 +112,37 @@ internal class StreamStatusesHandlerTest {
   }
 
   @Test
-  fun testListPerRunState() {
+  fun testListPerRunStateOriginal() {
     val connectionId = UUID.randomUUID()
     val apiReq = ConnectionIdRequestBody().connectionId(connectionId)
     val domainItem = StreamStatusBuilder().build()
     val apiItem = StreamStatusRead()
     val apiResp = StreamStatusReadList().streamStatuses(listOf<@Valid StreamStatusRead?>(apiItem))
 
+    // Use original implementation (feature flag disabled)
+    every { featureFlagClient.boolVariation(UseOptimizedStreamStatusQuery, Connection(connectionId)) } returns false
     every { repo.findAllPerRunStateByConnectionId(connectionId) } returns listOf(domainItem)
     every { mapper.map(domainItem) } returns apiItem
 
     Assertions.assertEquals(apiResp, handler.listStreamStatusPerRunState(apiReq))
   }
 
+  @Test
+  fun testListPerRunStateOptimized() {
+    val connectionId = UUID.randomUUID()
+    val apiReq = ConnectionIdRequestBody().connectionId(connectionId)
+    val domainItem = StreamStatusBuilder().build()
+    val apiItem = StreamStatusRead()
+    val apiResp = StreamStatusReadList().streamStatuses(listOf<@Valid StreamStatusRead?>(apiItem))
+
+    // Use optimized implementation (feature flag enabled)
+    every { featureFlagClient.boolVariation(UseOptimizedStreamStatusQuery, Connection(connectionId)) } returns true
+    every { repo.findAllPerRunStateByConnectionIdWithRecentJobsFilter(connectionId, 100) } returns listOf(domainItem)
+    every { mapper.map(domainItem) } returns apiItem
+
+    Assertions.assertEquals(apiResp, handler.listStreamStatusPerRunState(apiReq))
+  }
+
   @Test
   fun testUptimeHistory() {
     val connectionId = UUID.randomUUID()
@@ -379,6 +402,7 @@ internal class StreamStatusesHandlerTest {
         StreamStatusesMapper(),
         jobHistoryHandler,
         jobPersistence,
+        featureFlagClient,
       )
     Assertions.assertEquals(expected, handlerWithRealMapper.getConnectionUptimeHistory(apiReq))
   }