Merge pull request #217 from Sunbird-Lern/yugabyte_migration

Issue #SBCOSS-00: Replace QueryBuilder.batch() with BatchStatement for YugabyteDB compatibility

Author: sntiwari1

lms-jobs/activity-aggregate-updater/src/main/scala/org/sunbird/job/aggregate/functions/ActivityAggregatesFunction.scala

@@ -4,6 +4,7 @@ import java.lang.reflect.Type
 import java.util.concurrent.TimeUnit
 import com.datastax.driver.core.Row
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Select, Update}
+import com.datastax.driver.core.BatchStatement
 import com.google.gson.Gson
 import com.google.gson.reflect.TypeToken
 import com.twitter.storehaus.cache.TTLCache
@@ -270,18 +271,19 @@ class ActivityAggregatesFunction(config: ActivityAggregateUpdaterConfig, httpUti
   }
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
       } else {
-        val msg = "Database update has failed: " + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.info(msg)
         throw new Exception(msg)
       }

lms-jobs/activity-aggregate-updater/src/main/scala/org/sunbird/job/aggregate/functions/CollectionProgressCompleteFunction.scala

@@ -3,6 +3,7 @@ package org.sunbird.job.aggregate.functions
 import java.util.UUID
 
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Select, Update}
+import com.datastax.driver.core.BatchStatement
 import com.google.gson.Gson
 import org.apache.flink.api.common.typeinfo.TypeInformation
 import org.apache.flink.configuration.Configuration
@@ -103,21 +104,22 @@ class CollectionProgressCompleteFunction(config: ActivityAggregateUpdaterConfig)
   }
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
       logger.info("is cassandra cluster available =>"+(null !=cassandraUtil.session))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       logger.info("result after update => "+result)
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
         metrics.incCounter(config.enrolmentCompleteCount)
       } else {
-        val msg = "Database update has failed" + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.error(msg)
         throw new Exception(msg)
       }

lms-jobs/activity-aggregate-updater/src/main/scala/org/sunbird/job/aggregate/functions/CollectionProgressUpdateFunction.scala

@@ -1,6 +1,7 @@
 package org.sunbird.job.aggregate.functions
 
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Select, Update}
+import com.datastax.driver.core.BatchStatement
 import org.apache.flink.api.common.typeinfo.TypeInformation
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.streaming.api.functions.ProcessFunction
@@ -76,18 +77,19 @@ class CollectionProgressUpdateFunction(config: ActivityAggregateUpdaterConfig)(i
   }
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
       } else {
-        val msg = "Database update has failed" + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.error(msg)
         throw new Exception(msg)
       }

lms-jobs/enrolment-reconciliation/src/main/scala/org/sunbird/job/recounciliation/functions/EnrolmentReconciliationFn.scala

@@ -2,6 +2,7 @@ package org.sunbird.job.recounciliation.functions
 
 import com.datastax.driver.core.Row
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Select, Update}
+import com.datastax.driver.core.BatchStatement
 import com.google.gson.Gson
 import com.google.gson.reflect.TypeToken
 import com.twitter.storehaus.cache.TTLCache
@@ -298,18 +299,19 @@ class EnrolmentReconciliationFn(config: EnrolmentReconciliationConfig,  httpUtil
   }
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
       } else {
-        val msg = "Database update has failed: " + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.error(msg)
         throw new Exception(msg)
       }

lms-jobs/enrolment-reconciliation/src/main/scala/org/sunbird/job/recounciliation/functions/ProgressCompleteFunction.scala

@@ -1,6 +1,7 @@
 package org.sunbird.job.recounciliation.functions
 
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Select, Update}
+import com.datastax.driver.core.BatchStatement
 import com.google.gson.Gson
 import org.apache.flink.api.common.typeinfo.TypeInformation
 import org.apache.flink.configuration.Configuration
@@ -84,19 +85,20 @@ class ProgressCompleteFunction(config: EnrolmentReconciliationConfig)(implicit v
   }
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
         metrics.incCounter(config.enrolmentCompleteCount)
       } else {
-        val msg = "Database update has failed" + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.error(msg)
         throw new Exception(msg)
       }

lms-jobs/enrolment-reconciliation/src/main/scala/org/sunbird/job/recounciliation/functions/ProgressUpdateFunction.scala

@@ -1,6 +1,7 @@
 package org.sunbird.job.recounciliation.functions
 
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Select, Update}
+import com.datastax.driver.core.BatchStatement
 import org.apache.flink.api.common.typeinfo.TypeInformation
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.streaming.api.functions.ProcessFunction
@@ -65,18 +66,19 @@ class ProgressUpdateFunction(config: EnrolmentReconciliationConfig)(implicit val
   }
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
       } else {
-        val msg = "Database update has failed" + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.error(msg)
         throw new Exception(msg)
       }

user-org-jobs/user-ownership-transfer/src/main/scala/org/sunbird/job/ownershiptransfer/functions/UserOwnershipTransferFunction.scala

@@ -1,6 +1,7 @@
 package org.sunbird.job.ownershiptransfer.functions
 
 import com.datastax.driver.core.querybuilder.{QueryBuilder, Update}
+import com.datastax.driver.core.BatchStatement
 import org.apache.commons.lang3.StringUtils
 import org.apache.flink.api.common.typeinfo.TypeInformation
 import org.apache.flink.configuration.Configuration
@@ -130,19 +131,20 @@ class UserOwnershipTransferFunction(config: UserOwnershipTransferConfig, httpUti
 
 
   /**
-   * Method to update the specific table in a batch format.
+   * Method to update the specific table using YugabyteDB batch operations.
+   * Uses BatchStatement and cassandraUtil.update() for proper execution.
    */
   def updateDB(batchSize: Int, queriesList: List[Update.Where])(implicit metrics: Metrics): Unit = {
     val groupedQueries = queriesList.grouped(batchSize).toList
     groupedQueries.foreach(queries => {
-      val cqlBatch = QueryBuilder.batch()
-      queries.map(query => cqlBatch.add(query))
-      val result = cassandraUtil.upsert(cqlBatch.toString)
+      val batch = new BatchStatement()
+      queries.foreach(query => batch.add(query))
+      val result = cassandraUtil.update(batch)  // Use update() instead of upsert()
       if (result) {
         metrics.incCounter(config.dbUpdateCount)
         logger.info("DB update successful")
       } else {
-        val msg = "Database update has failed: " + cqlBatch.toString
+        val msg = "Database update has failed: " + batch.toString
         logger.info(msg)
         throw new Exception(msg)
       }