Acdc count metrics on DB table records (#52)

Author: schowsf

acdc-core/src/main/scala/com/salesforce/mce/acdc/db/DatasetInstanceQuery.scala

@@ -35,6 +35,8 @@ object DatasetInstanceQuery {
 
   def forDataset(dataset: String) = DatasetInstanceTable().filter(_.dataset === dataset).result
 
+  def count() = DatasetInstanceTable().length.result
+
   def filter(
     dataset: Option[String],
     like: Option[String],
@@ -59,7 +61,8 @@ object DatasetInstanceQuery {
         (t, d) => t.filter(r => r.dataset === d)
       )
 
-    like.foldLeft(filteredByDataset)((t, l) => t.filter(r => r.name.like(l)))
+    like
+      .foldLeft(filteredByDataset)((t, l) => t.filter(r => r.name.like(l)))
       .sortBy(sortByColumn)
       .drop(offset)
       .take(limit)

acdc-core/src/main/scala/com/salesforce/mce/acdc/db/DatasetLineageQuery.scala

@@ -13,6 +13,8 @@ import slick.jdbc.PostgresProfile.api._
 
 object DatasetLineageQuery {
 
+  def count() = DatasetLineageTable().length.result
+
   case class ForDestination(dest: String) {
 
     def table = DatasetLineageTable().filter(_.toDataset === dest)

acdc-core/src/main/scala/com/salesforce/mce/acdc/db/DatasetQuery.scala

@@ -29,6 +29,8 @@ object DatasetQuery {
 
   def now() = LocalDateTime.now()
 
+  def count() = DatasetTable().length.result
+
   case class ForName(name: String) {
 
     def table = DatasetTable().filter(_.name === name)

acdc-ws/app/Module.scala

@@ -8,7 +8,7 @@
 import com.google.inject.AbstractModule
 
 import services.{Metric, PrometheusMetric}
-import tasks.{AuthSettingReloadTask, DataInstExpirationTask}
+import tasks.{AuthSettingReloadTask, DataCountTask, DataInstExpirationTask}
 import utils.{Authorization, AuthorizationSettings}
 
 class Module extends AbstractModule {
@@ -22,6 +22,8 @@ class Module extends AbstractModule {
     bind(classOf[Metric]).to(classOf[PrometheusMetric])
     // Data forget task
     bind(classOf[DataInstExpirationTask]).asEagerSingleton()
+    // Data count task
+    bind(classOf[DataCountTask]).asEagerSingleton()
   }
 
 }

acdc-ws/app/services/Metric.scala

@@ -54,6 +54,8 @@ trait Metric {
 
   def clear(): Unit
 
+  def countDB(table: String, cnt: Double): Unit
+
 }
 
 @Singleton
@@ -77,6 +79,17 @@ class PrometheusMetric @Inject() (implicit ec: ExecutionContext) extends Metric
     .quantile(0.99d, 0.001d)
     .register
 
+  private val dbCountGauge = Gauge
+    .build()
+    .name("acdc_db_record_count")
+    .labelNames("table")
+    .help("acdc DB table records count")
+    .register
+
+  override def countDB(table: String, cnt: Double): Unit = {
+    dbCountGauge.labels(table).set(cnt)
+  }
+
   override def incrementStatusCount(status: String): Unit = {
     httpStatusCount.labels(status).inc()
   }

acdc-ws/app/tasks/DataCountTask.scala

@@ -0,0 +1,55 @@
+package tasks
+
+import javax.inject.Inject
+
+import scala.concurrent.ExecutionContext
+import scala.concurrent.duration.DurationInt
+
+import akka.actor.ActorSystem
+import play.api.Logging
+
+import com.salesforce.mce.acdc.db.{
+  AcdcDatabase,
+  DatasetInstanceQuery,
+  DatasetLineageQuery,
+  DatasetQuery
+}
+import services.{DatabaseService, Metric}
+import utils.DbConfig
+
+class DataCountTask @Inject() (
+  actorSystem: ActorSystem,
+  dbService: DatabaseService,
+  dbConfig: DbConfig,
+  metric: Metric
+)(implicit
+  ec: ExecutionContext
+) extends Logging {
+
+  val countFrequency: Int = dbConfig.countTaskFrequencyMinute
+  def db: AcdcDatabase = dbService.db
+
+  def refresh(): Unit = {
+    actorSystem.scheduler.scheduleOnce(countFrequency.minutes) {
+      val datasetCount = db.sync(DatasetQuery.count())
+      val datasetInstanceCount = db.sync(DatasetInstanceQuery.count())
+      val datasetLineageCount = db.sync(DatasetLineageQuery.count())
+
+      logger.debug(s"Counted $datasetCount records from dataset ...")
+      logger.debug(s"Counted $datasetInstanceCount records from dataset_instance ...")
+      logger.debug(s"Counted $datasetLineageCount records from dataset_lineage ...")
+
+      metric.countDB("dataset", datasetCount.toDouble)
+      metric.countDB("dataset-instance", datasetInstanceCount.toDouble)
+      metric.countDB("dataset-lineage", datasetLineageCount.toDouble)
+
+      refresh()
+    }
+  }
+
+  if (countFrequency > 0) {
+    refresh()
+  } else {
+    logger.info("data count frequency is not a positive integer, data count task skipped.")
+  }
+}

acdc-ws/app/utils/DbConfig.scala

@@ -11,8 +11,9 @@ import scala.util.{Failure, Success, Try}
 
 import com.typesafe.config.ConfigException
 import com.typesafe.config.ConfigFactory
+import play.api.Logging
 
-class DbConfig() {
+class DbConfig() extends Logging {
 
   val config = ConfigFactory.load().getConfig("acdc.db")
 
@@ -22,4 +23,12 @@ class DbConfig() {
     case Failure(e) => throw e
   }
 
+  def countTaskFrequencyMinute: Int = Try(config.getInt(s"count-task-frequency-minute")) match {
+    case Success(d) => d
+    case Failure(e: ConfigException.Missing) =>
+      logger.warn("No config found for count-task-frequency-minute, defaulting to 0.")
+      0
+    case Failure(e) => throw e
+  }
+
 }

acdc-ws/conf/application.conf

@@ -32,6 +32,11 @@ acdc.auth = {
 acdc.db {
     # maximum length in days to keep
     ttl = ${?DATA_TTL}
+
+    # count and report records in dataset, datasetInstance, and datasetLineage tables
+    # 0 is to skip data count task, set it to a positive Int for effective count frequency
+    count-task-frequency-minute = 0
+    count-task-frequency-minute = ${?COUNT_TASK_FREQUENCY_MINUTE}
 }
 
 ## Akka

version.sbt

@@ -1 +1 @@
- ThisBuild / version := "0.8.11"
+ ThisBuild / version := "0.8.12"