float up gb fetcher methods

Author: nikhil-zlai

online/src/main/scala/ai/chronon/online/fetcher/Fetcher.scala

@@ -123,7 +123,7 @@ class Fetcher(val kvStore: KVStore,
   }
 
   def fetchGroupBys(requests: Seq[Request]): Future[Seq[Response]] = {
-    joinPartFetcher.groupByFetcher.fetchGroupBys(requests)
+    joinPartFetcher.fetchGroupBys(requests)
   }
 
   def fetchJoin(requests: Seq[Request], joinConf: Option[api.Join] = None): Future[Seq[Response]] = {

online/src/main/scala/ai/chronon/online/fetcher/JoinPartFetcher.scala

@@ -19,7 +19,8 @@ package ai.chronon.online.fetcher
 import ai.chronon.api.Extensions._
 import ai.chronon.api._
 import ai.chronon.online._
-import ai.chronon.online.fetcher.Fetcher.{PrefixedRequest, Request, Response}
+import ai.chronon.online.fetcher.Fetcher.{ColumnSpec, PrefixedRequest, Request, Response}
+import ai.chronon.online.fetcher.FetcherCache.BatchResponses
 import org.slf4j.{Logger, LoggerFactory}
 
 import scala.collection.Seq
@@ -30,9 +31,28 @@ class JoinPartFetcher(fetchContext: FetchContext, metadataStore: MetadataStore)
 
   @transient implicit lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
-  private[online] val groupByFetcher = new GroupByFetcher(fetchContext, metadataStore)
+  private val groupByFetcher = new GroupByFetcher(fetchContext, metadataStore)
   private implicit val executionContext: ExecutionContext = fetchContext.getOrCreateExecutionContext
 
+  def fetchGroupBys(requests: Seq[Request]): Future[Seq[Response]] = {
+    groupByFetcher.fetchGroupBys(requests)
+  }
+
+  // ----- START -----
+  // floated up to makes tests easy
+  def fetchColumns(specs: Seq[ColumnSpec]): Future[Map[ColumnSpec, Response]] = {
+    groupByFetcher.fetchColumns(specs)
+  }
+
+  def getServingInfo(existing: GroupByServingInfoParsed, batchResponses: BatchResponses): GroupByServingInfoParsed = {
+    groupByFetcher.getServingInfo(existing, batchResponses)
+  }
+
+  def isCacheSizeConfigured: Boolean = {
+    groupByFetcher.isCacheSizeConfigured
+  }
+  // ---- END ----
+
   // prioritize passed in joinOverrides over the ones in metadata store
   // used in stream-enrichment and in staging testing
   def fetchJoins(requests: Seq[Request], joinConf: Option[Join] = None): Future[Seq[Response]] = {

online/src/test/scala/ai/chronon/online/test/FetcherBaseTest.scala

@@ -85,16 +85,16 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
         val response = Response(request, Success(Map(request.name -> "100")))
         Future.successful(Seq(response))
       }
-    }).when(fetcherBase).groupByFetcher.fetchGroupBys(any())
+    }).when(fetcherBase).fetchGroupBys(any())
 
     // Map should contain query with valid response
-    val queryResults = Await.result(fetcherBase.groupByFetcher.fetchColumns(Seq(query)), 1.second)
+    val queryResults = Await.result(fetcherBase.fetchColumns(Seq(query)), 1.second)
     queryResults.contains(query) shouldBe true
     queryResults.get(query).map(_.values) shouldBe Some(Success(Map(s"$GroupBy.$Column" -> "100")))
 
     // GroupBy request sent to KV store for the query
     val requestsCaptor = ArgumentCaptor.forClass(classOf[Seq[_]])
-    verify(fetcherBase, times(1)).groupByFetcher.fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
+    verify(fetcherBase, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
     val actualRequest = requestsCaptor.getValue.asInstanceOf[Seq[Request]].headOption
     actualRequest shouldNot be(None)
     actualRequest.get.name shouldBe s"${query.groupByName}.${query.columnName}"
@@ -114,18 +114,18 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
         val responses = requests.map(r => Response(r, Success(Map(r.name -> "100"))))
         Future.successful(responses)
       }
-    }).when(fetcherBase).groupByFetcher.fetchGroupBys(any())
+    }).when(fetcherBase).fetchGroupBys(any())
 
     // Map should contain query with valid response
-    val queryResults = Await.result(fetcherBase.groupByFetcher.fetchColumns(Seq(guestQuery, hostQuery)), 1.second)
+    val queryResults = Await.result(fetcherBase.fetchColumns(Seq(guestQuery, hostQuery)), 1.second)
     queryResults.contains(guestQuery) shouldBe true
     queryResults.get(guestQuery).map(_.values) shouldBe Some(Success(Map(s"${GuestKey}_$GroupBy.$Column" -> "100")))
     queryResults.contains(hostQuery) shouldBe true
     queryResults.get(hostQuery).map(_.values) shouldBe Some(Success(Map(s"${HostKey}_$GroupBy.$Column" -> "100")))
 
     // GroupBy request sent to KV store for the query
     val requestsCaptor = ArgumentCaptor.forClass(classOf[Seq[_]])
-    verify(fetcherBase, times(1)).groupByFetcher.fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
+    verify(fetcherBase, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
     val actualRequests = requestsCaptor.getValue.asInstanceOf[Seq[Request]]
     actualRequests.length shouldBe 2
     actualRequests.head.name shouldBe s"${guestQuery.groupByName}.${guestQuery.columnName}"
@@ -143,10 +143,10 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
       def answer(invocation: InvocationOnMock): Future[Seq[Response]] = {
         Future.successful(Seq())
       }
-    }).when(fetcherBase).groupByFetcher.fetchGroupBys(any())
+    }).when(fetcherBase).fetchGroupBys(any())
 
     // Map should contain query with Failure response
-    val queryResults = Await.result(fetcherBase.groupByFetcher.fetchColumns(Seq(query)), 1.second)
+    val queryResults = Await.result(fetcherBase.fetchColumns(Seq(query)), 1.second)
     queryResults.contains(query) shouldBe true
     queryResults.get(query).map(_.values) match {
       case Some(Failure(_: IllegalStateException)) => succeed
@@ -155,7 +155,7 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
 
     // GroupBy request sent to KV store for the query
     val requestsCaptor = ArgumentCaptor.forClass(classOf[Seq[_]])
-    verify(fetcherBase, times(1)).groupByFetcher.fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
+    verify(fetcherBase, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
     val actualRequest = requestsCaptor.getValue.asInstanceOf[Seq[Request]].headOption
     actualRequest shouldNot be(None)
     actualRequest.get.name shouldBe query.groupByName + "." + query.columnName
@@ -166,16 +166,16 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
   it should "get serving info should call update serving info if batch response is from kv store" in {
     val oldServingInfo = mock[GroupByServingInfoParsed]
     val updatedServingInfo = mock[GroupByServingInfoParsed]
-    doReturn(updatedServingInfo).when(fetcherBase).groupByFetcher.getServingInfo(any(), any())
+    doReturn(updatedServingInfo).when(fetcherBase).getServingInfo(any(), any())
 
     val batchTimedValuesSuccess = Success(Seq(TimedValue(Array(1.toByte), 2000L)))
     val kvStoreBatchResponses = BatchResponses(batchTimedValuesSuccess)
 
-    val result = fetcherBase.groupByFetcher.getServingInfo(oldServingInfo, kvStoreBatchResponses)
+    val result = fetcherBase.getServingInfo(oldServingInfo, kvStoreBatchResponses)
 
     // updateServingInfo is called
     result shouldEqual updatedServingInfo
-    verify(fetcherBase).groupByFetcher.getServingInfo(any(), any())
+    verify(fetcherBase).getServingInfo(any(), any())
   }
 
   // If a batch response is cached, the serving info should be refreshed. This is needed to prevent
@@ -194,12 +194,12 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
     doReturn(groupByOpsMock).when(oldServingInfo).groupByOps
 
     val cachedBatchResponses = BatchResponses(mock[FinalBatchIr])
-    val result = fetcherBase.groupByFetcher.getServingInfo(oldServingInfo, cachedBatchResponses)
+    val result = fetcherBase.getServingInfo(oldServingInfo, cachedBatchResponses)
 
     // FetcherBase.updateServingInfo is not called, but getGroupByServingInfo.refresh() is.
     result shouldEqual oldServingInfo
     verify(ttlCache).refresh(any())
-    verify(fetcherBase, never()).groupByFetcher.getServingInfo(any(), any())
+    verify(fetcherBase, never()).getServingInfo(any(), any())
   }
 
   it should "is caching enabled correctly determine if cache is enabled" in {
@@ -224,7 +224,7 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
 
     val fetcherBaseWithFlagStore =
       spy[fetcher.JoinPartFetcher](new fetcher.JoinPartFetcher(fetchContext, new MetadataStore(fetchContext)))
-    when(fetcherBaseWithFlagStore.groupByFetcher.isCacheSizeConfigured).thenReturn(true)
+    when(fetcherBaseWithFlagStore.isCacheSizeConfigured).thenReturn(true)
 
     def buildGroupByWithCustomJson(name: String): GroupBy = Builders.GroupBy(metaData = Builders.MetaData(name = name))