Merge pull request #392 from revinjchalil/2.4

Fix streaming checkpoint issues

Author: FabianMeiswinkel

pom.xml

@@ -21,7 +21,7 @@ limitations under the License.
     <groupId>com.microsoft.azure</groupId>
     <artifactId>azure-cosmosdb-spark_2.4.0_2.11</artifactId>
     <packaging>jar</packaging>
-    <version>3.0.5</version>
+    <version>3.0.6</version>
     <name>${project.groupId}:${project.artifactId}</name>
     <description>Spark Connector for Microsoft Azure CosmosDB</description>
     <url>http://azure.microsoft.com/en-us/services/documentdb/</url>

src/main/scala/com/microsoft/azure/cosmosdb/spark/Constants.scala

@@ -23,6 +23,6 @@
 package com.microsoft.azure.cosmosdb.spark
 
 object Constants {
-  val currentVersion = "2.4.0_2.11-3.0.5"
+  val currentVersion = "2.4.0_2.11-3.0.6"
   val userAgentSuffix = s" SparkConnector/$currentVersion"
 }

src/main/scala/com/microsoft/azure/cosmosdb/spark/CosmosDBConnection.scala

@@ -230,6 +230,11 @@ private[spark] case class CosmosDBConnection(config: Config) extends CosmosDBLog
         }
       }
 
+      // set nextContinuation if the given partition does not have an existing checkpoint file and the "startFromBeginning" is False
+      if (nextContinuation == null || nextContinuation.isEmpty){
+        nextContinuation = feedResponse.getResponseContinuation
+      }
+
       logDebug(s"<-- readChangeFeed, Count: ${cfDocuments.length.toString}, NextContinuation: $nextContinuation")
 
       updateTokenFunc(originalContinuation, nextContinuation, partitionId)

src/main/scala/com/microsoft/azure/cosmosdb/spark/rdd/CosmosDBRDDIterator.scala

@@ -50,11 +50,11 @@ object CosmosDBRDDIterator {
   var lastFeedOptions: FeedOptions = _
   var hdfsUtils: HdfsUtils = _
 
-  def initializeHdfsUtils(hadoopConfig: Map[String, String]): Any = {
+  def initializeHdfsUtils(hadoopConfig: Map[String, String], changeFeedCheckpointLocation: String): Any = {
     if (hdfsUtils == null) {
       this.synchronized {
         if (hdfsUtils == null) {
-          hdfsUtils = HdfsUtils(hadoopConfig)
+          hdfsUtils = HdfsUtils(hadoopConfig, changeFeedCheckpointLocation)
         }
       }
     }
@@ -133,7 +133,11 @@ class CosmosDBRDDIterator(hadoopConfig: mutable.Map[String, String],
   extends Iterator[Document]
     with CosmosDBLoggingTrait {
 
-  CosmosDBRDDIterator.initializeHdfsUtils(hadoopConfig.toMap)
+  val changeFeedCheckpointLocation: String = config
+    .get[String](CosmosDBConfig.ChangeFeedCheckpointLocation)
+    .getOrElse(StringUtils.EMPTY)
+
+  CosmosDBRDDIterator.initializeHdfsUtils(hadoopConfig.toMap, changeFeedCheckpointLocation)
 
   // The continuation token for the target CosmosDB partition
   private var cfCurrentToken: String = _
@@ -223,9 +227,6 @@ class CosmosDBRDDIterator(hadoopConfig: mutable.Map[String, String],
 
       val objectMapper: ObjectMapper = new ObjectMapper()
 
-      val changeFeedCheckpointLocation: String = config
-        .get[String](CosmosDBConfig.ChangeFeedCheckpointLocation)
-        .getOrElse(StringUtils.EMPTY)
       val queryName: String = config
         .get[String](CosmosDBConfig.ChangeFeedQueryName)
         .get

src/main/scala/com/microsoft/azure/cosmosdb/spark/streaming/CosmosDBSource.scala

@@ -63,13 +63,13 @@ private[spark] class CosmosDBSource(sqlContext: SQLContext,
     }
 
     if (currentSchema == null) {
+      val changeFeedCheckpointLocation: String = streamConfigMap
+        .getOrElse(CosmosDBConfig.ChangeFeedCheckpointLocation, StringUtils.EMPTY)
       CosmosDBRDDIterator.initializeHdfsUtils(HdfsUtils.getConfigurationMap(
-        sqlContext.sparkSession.sparkContext.hadoopConfiguration).toMap)
+        sqlContext.sparkSession.sparkContext.hadoopConfiguration).toMap, changeFeedCheckpointLocation)
 
       // Delete current tokens and next tokens checkpoint directories to ensure change feed starts from beginning if set
       if (streamConfigMap.getOrElse(CosmosDBConfig.ChangeFeedStartFromTheBeginning, String.valueOf(false)).toBoolean) {
-        val changeFeedCheckpointLocation: String = streamConfigMap
-          .getOrElse(CosmosDBConfig.ChangeFeedCheckpointLocation, StringUtils.EMPTY)
         val queryName = Config(streamConfigMap)
           .get[String](CosmosDBConfig.ChangeFeedQueryName).get
         val currentTokensCheckpointPath = changeFeedCheckpointLocation + "/" + HdfsUtils.filterFilename(queryName)

src/main/scala/com/microsoft/azure/cosmosdb/spark/streaming/DefaultCosmosDBWriteStreamPoisonMessageNotificationHandler.scala

@@ -22,16 +22,17 @@
   */
 package com.microsoft.azure.cosmosdb.spark.streaming
 
-import com.microsoft.azure.cosmosdb.{Document, ResourceResponse, RequestOptions}
+import com.microsoft.azure.cosmosdb.{Document, RequestOptions, ResourceResponse}
 import com.microsoft.azure.cosmosdb.spark.config.{Config, CosmosDBConfig}
 import com.microsoft.azure.cosmosdb.spark.CosmosDBLoggingTrait
 import com.microsoft.azure.cosmosdb.spark.util.HdfsUtils
-
 import java.io.{FileNotFoundException, PrintWriter, StringWriter}
 import java.time.Instant
 import java.time.format.DateTimeFormatter
 import java.util.concurrent.atomic.AtomicLong
 
+import org.apache.commons.lang3.StringUtils
+
 class DefaultCosmosDBWriteStreamPoisonMessageNotificationHandler(configMap: Map[String, String])
     extends CosmosDBWriteStreamPoisonMessageNotificationHandler
     with CosmosDBLoggingTrait
@@ -45,9 +46,9 @@ class DefaultCosmosDBWriteStreamPoisonMessageNotificationHandler(configMap: Map[
         .getOrElse(
             CosmosDBConfig.PoisonMessageLocation,
             String.valueOf(CosmosDBConfig.DefaultPoisonMessageLocation))
-    
+
     private lazy val hdfsUtils: HdfsUtils = {
-        HdfsUtils(configMap)
+        HdfsUtils(configMap, poisonMessageLocation)
     }
 
     def onPoisonMessage(lastError: Throwable, document: Document): Unit =
@@ -61,7 +62,7 @@ class DefaultCosmosDBWriteStreamPoisonMessageNotificationHandler(configMap: Map[
 
         val payload = document.toJson()
 
-        logWarning(s"POSION MESSAGE Id: $id, Error: $error, Document payload: $payload")
+        logWarning(s"POISON MESSAGE Id: $id, Error: $error, Document payload: $payload")
 
         if (this.poisonMessageLocation != "")
         {

src/main/scala/com/microsoft/azure/cosmosdb/spark/util/HdfsUtils.scala

@@ -31,16 +31,17 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, LocatedFileStatus, Path, RemoteIterator}
 
 import scala.collection.mutable
+import java.net.URI
 
-case class HdfsUtils(configMap: Map[String, String]) extends CosmosDBLoggingTrait {
+case class HdfsUtils(configMap: Map[String, String], changeFeedCheckpointLocation: String) extends CosmosDBLoggingTrait {
   private val fsConfig: Configuration = {
     val config = new Configuration()
     configMap.foreach(e => config.set(e._1, e._2))
     config
   }
 
   private val maxRetryCount = 10
-  private val fs = FileSystem.get(fsConfig)
+  private val fs = FileSystem.get(new URI(changeFeedCheckpointLocation), fsConfig)
 
   def write(base: String, filePath: String, content: String): Unit = {
     val path = new Path(base + "/" + filePath)