Feature/add kusto streaming feature

.github/workflows/release.yml

@@ -31,8 +31,11 @@ jobs:
           access_token=$(az account get-access-token --resource=${{ secrets.APP_ID }} --scope=${{ secrets.CLUSTER }}/.default --query accessToken -o tsv)
           echo "ACCESS_TOKEN=$access_token" >> $GITHUB_ENV
       - name: Run the Maven verify phase
+        run: mvn clean verify -DkustoAadAppId=${{ secrets.APP_ID }} -DkustoAadAuthorityID=${{ secrets.TENANT_ID }} -DkustoAadAppSecret=${{ secrets.APP_SECRET }} -DkustoDatabase=${{ secrets.DATABASE }} -DkustoCluster=${{ secrets.CLUSTER }}
         env:
+          kustoAadAppId: ${{ secrets.APP_ID }}
           kustoAadAuthorityID: ${{ secrets.TENANT_ID }}
+          kustoAadAppSecret: ${{ secrets.APP_SECRET }}
           kustoDatabase: ${{ secrets.DATABASE }}
           kustoCluster: ${{ secrets.CLUSTER }}
           kustoAadAppId: ${{secrets.APP_ID}}

connector/pom.xml

@@ -673,6 +673,13 @@
                     </transformationSets>
                 </configuration>
             </plugin>
+            <plugin>
+                <groupId>org.antipathy</groupId>
+                <artifactId>mvn-scalafmt_${scala.version.major}</artifactId>
+                <configuration>
+                    <configLocation>${project.parent.basedir}/.scalafmt.conf</configLocation> <!-- path to config -->
+                </configuration>
+            </plugin>
         </plugins>
     </build>
     <repositories>

connector/src/main/scala/com/microsoft/kusto/spark/common/KustoOptions.scala

@@ -84,15 +84,16 @@ case class KustoCoordinates(
     ingestionUrl: Option[String] = None)
 
 /**
- * *******************************************************************************
+ * *********************************************************************************
  */
-/*                                    NOTE!!!                                       */
-/* These options are intended for testing, experimentation and debug.               */
-/* They may not be used in a production environment                                 */
-/* Interface stability is not guaranteed: options may be removed or changed freely  */
+//                                    NOTE!!!
+// These options are intended for testing, experimentation and debug.
+// They may not be used in a production environment
+// Interface stability is not guaranteed: options may be removed or changed freely
 /**
- * *******************************************************************************
+ * **********************************************************************************
  */
+
 private[kusto] object KustoDebugOptions {
   private val kustoOptionNames = collection.mutable.Set[String]()
 

connector/src/main/scala/com/microsoft/kusto/spark/datasink/FinalizeHelper.scala

@@ -5,8 +5,7 @@ package com.microsoft.kusto.spark.datasink
 
 import com.azure.data.tables.implementation.models.TableServiceErrorException
 import com.fasterxml.jackson.databind.ObjectMapper
-
-import java.util.concurrent.TimeUnit
+import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule
 import com.microsoft.azure.kusto.data.ClientRequestProperties
 import com.microsoft.azure.kusto.ingest.result.{IngestionStatus, OperationStatus}
 import com.microsoft.kusto.spark.authentication.KustoAuthentication
@@ -20,20 +19,20 @@ import com.microsoft.kusto.spark.utils.KustoConstants.IngestSkippedTrace
 import com.microsoft.kusto.spark.utils.{
   ExtendedKustoClient,
   KustoClientCache,
-  KustoConstants,
   KustoDataSourceUtils => KDSU
 }
 import org.apache.spark.SparkContext
 import org.apache.spark.util.CollectionAccumulator
 
 import java.time.Instant
+import java.util.concurrent.TimeUnit
 import scala.collection.JavaConverters._
-import scala.concurrent.{Await, Future, TimeoutException}
 import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.{Await, Future, TimeoutException}
 
 object FinalizeHelper {
   private val myName = this.getClass.getSimpleName
-
+  private val mapper = new ObjectMapper().registerModule(new JavaTimeModule())
   private[kusto] def finalizeIngestionWhenWorkersSucceeded(
       coordinates: KustoCoordinates,
       batchIdIfExists: String,
@@ -48,19 +47,18 @@ object FinalizeHelper {
       sinkStartTime: Instant): Unit = {
     if (!kustoClient.shouldIngestData(
         coordinates,
-        writeOptions.ingestionProperties,
+        writeOptions.maybeSparkIngestionProperties,
         tableExists,
         crp)) {
       KDSU.logInfo(myName, s"$IngestSkippedTrace '${coordinates.table}'")
     } else {
       val mergeTask = Future {
-        val loggerName = myName
         val requestId = writeOptions.requestId
         val ingestionInfoString =
           s"RequestId: $requestId cluster: '${coordinates.clusterAlias}', " +
             s"database: '${coordinates.database}', table: '$tmpTableName' $batchIdIfExists"
         KDSU.logInfo(
-          loggerName,
+          myName,
           s"Polling on ingestion results for requestId: $requestId, will move data to " +
             s"destination table when finished")
 
@@ -69,7 +67,6 @@ object FinalizeHelper {
             partitionsResults.value.asScala.foreach(partitionResult =>
               pollOnResult(
                 partitionResult,
-                loggerName,
                 requestId,
                 writeOptions.timeout.toMillis,
                 ingestionInfoString,
@@ -86,7 +83,6 @@ object FinalizeHelper {
               results.foreach(partitionResult =>
                 pollOnResult(
                   partitionResult,
-                  loggerName,
                   requestId,
                   writeOptions.timeout.toMillis,
                   ingestionInfoString,
@@ -188,7 +184,6 @@ object FinalizeHelper {
 
   def pollOnResult(
       partitionResult: PartitionResult,
-      loggerName: String,
       requestId: String,
       timeout: Long,
       ingestionInfoString: String,
@@ -201,20 +196,16 @@ object FinalizeHelper {
             finalRes = Some(partitionResult.ingestionResult.getIngestionStatusCollection.get(0))
             finalRes
           } catch {
-//          case e: RequestFailedException =>
-//            KDSU.logWarn(loggerName, "Failed to fetch operation status transiently - will keep polling. " +
-//              s"RequestId: $requestId. Error: ${ExceptionUtils.getStackTrace(e)}")
-//            None
             case e: TableServiceErrorException =>
               KDSU.reportExceptionAndThrow(
-                loggerName,
+                myName,
                 e,
                 s"TableServiceErrorException : RequestId: $requestId",
                 shouldNotThrow = true)
               None
             case e: Exception =>
               KDSU.reportExceptionAndThrow(
-                loggerName,
+                myName,
                 e,
                 s"Failed to fetch operation status. RequestId: $requestId")
               None
@@ -226,7 +217,7 @@ object FinalizeHelper {
           val pending = res.isDefined && res.get.status == OperationStatus.Pending
           if (pending) {
             KDSU.logDebug(
-              loggerName,
+              myName,
               s"Polling on result for partition: '${partitionResult.partitionId}' in requestId: $requestId, status is-'Pending'")
           }
           pending
@@ -235,48 +226,58 @@ object FinalizeHelper {
         maxWaitTimeBetweenCallsMillis = KDSU.WriteInitialMaxWaitTime.toMillis.toInt,
         maxWaitTimeAfterMinute = KDSU.WriteMaxWaitTime.toMillis.toInt)
       .await(timeout, TimeUnit.MILLISECONDS)
+    finalRes match {
+      case Some(ingestResults) =>
+        processIngestionStatusResults(
+          partitionResult.partitionId,
+          ingestionInfoString,
+          shouldThrowOnTagsAlreadyExists,
+          ingestResults)
+      case None => throw new RuntimeException("Failed to poll on ingestion status.")
+    }
+  }
 
-    if (finalRes.isDefined) {
-      finalRes.get.status match {
-        case OperationStatus.Pending =>
+  def processIngestionStatusResults(
+      partitionId: Int = 0,
+      ingestionInfoString: String,
+      shouldThrowOnTagsAlreadyExists: Boolean,
+      ingestionStatusResult: IngestionStatus): Unit = {
+    ingestionStatusResult.status match {
+      case OperationStatus.Pending =>
+        throw new RuntimeException(
+          s"Ingestion to Kusto failed on timeout failure. $ingestionInfoString,  partition: '$partitionId'")
+      case OperationStatus.Succeeded =>
+        KDSU.logInfo(
+          myName,
+          s"Ingestion to Kusto succeeded. $ingestionInfoString,  partition: '$partitionId', " +
+            s"from: '${ingestionStatusResult.ingestionSourcePath}' , Operation ${ingestionStatusResult.operationId}")
+      case OperationStatus.Skipped =>
+        // TODO: should we throw ?
+        KDSU.logInfo(
+          myName,
+          s"Ingestion to Kusto skipped. $ingestionInfoString, " +
+            s"partition: '$partitionId', from: '${ingestionStatusResult.ingestionSourcePath}', " +
+            s"Operation ${ingestionStatusResult.operationId}")
+      case otherStatus: Any =>
+        // TODO error code should be added to java client
+        if (ingestionStatusResult.errorCodeString != "Skipped_IngestByTagAlreadyExists") {
           throw new RuntimeException(
-            s"Ingestion to Kusto failed on timeout failure. $ingestionInfoString, " +
-              s"partition: '${partitionResult.partitionId}'")
-        case OperationStatus.Succeeded =>
-          KDSU.logInfo(
-            loggerName,
-            s"Ingestion to Kusto succeeded. $ingestionInfoString, " +
-              s"partition: '${partitionResult.partitionId}', from: '${finalRes.get.ingestionSourcePath}', " +
-              s"Operation ${finalRes.get.operationId}")
-        case OperationStatus.Skipped =>
-          // TODO: should we throw ?
-          KDSU.logInfo(
-            loggerName,
-            s"Ingestion to Kusto skipped. $ingestionInfoString, " +
-              s"partition: '${partitionResult.partitionId}', from: '${finalRes.get.ingestionSourcePath}', " +
-              s"Operation ${finalRes.get.operationId}")
-        case otherStatus: Any =>
-          // TODO error code should be added to java client
-          if (finalRes.get.errorCodeString != "Skipped_IngestByTagAlreadyExists") {
-            throw new RuntimeException(s"Ingestion to Kusto failed with status '$otherStatus'." +
-              s" $ingestionInfoString, partition: '${partitionResult.partitionId}'. Ingestion info: '${new ObjectMapper().writerWithDefaultPrettyPrinter
-                  .writeValueAsString(finalRes.get)}'")
-          } else if (shouldThrowOnTagsAlreadyExists) {
-            // TODO - think about this logic and other cases that should not throw all (maybe everything that starts with skip? this actualy
-            //  seems like a bug in engine that the operation status is not Skipped)
-            //  (Skipped_IngestByTagAlreadyExists is relevant for dedup flow only as in other cases we cancel the ingestion altogether)
-            throw new RuntimeException(s"Ingestion to Kusto skipped with status '$otherStatus'." +
-              s" $ingestionInfoString, partition: '${partitionResult.partitionId}'. Ingestion info: '${new ObjectMapper().writerWithDefaultPrettyPrinter
-                  .writeValueAsString(finalRes.get)}'")
-          }
-          KDSU.logInfo(
-            loggerName,
-            s"Ingestion to Kusto failed. $ingestionInfoString, " +
-              s"partition: '${partitionResult.partitionId}', from: '${finalRes.get.ingestionSourcePath}', " +
-              s"Operation ${finalRes.get.operationId}")
-      }
-    } else {
-      throw new RuntimeException("Failed to poll on ingestion status.")
+            s"Ingestion to Kusto failed with status '$otherStatus'." +
+              s" $ingestionInfoString, partition: '$partitionId'. Ingestion info: '${mapper.writerWithDefaultPrettyPrinter
+                  .writeValueAsString(ingestionStatusResult)}'")
+        } else if (shouldThrowOnTagsAlreadyExists) {
+          // TODO - think about this logic and other cases that should not throw all (maybe everything that starts with skip? this actualy
+          //  seems like a bug in engine that the operation status is not Skipped)
+          //  (Skipped_IngestByTagAlreadyExists is relevant for dedup flow only as in other cases we cancel the ingestion altogether)
+          throw new RuntimeException(s"Ingestion to Kusto skipped with status '$otherStatus'." +
+            s" $ingestionInfoString, partition: '$partitionId'. Ingestion info: '${new ObjectMapper().writerWithDefaultPrettyPrinter
+                .writeValueAsString(ingestionStatusResult)}'")
+        }
+        KDSU.logInfo(
+          myName,
+          s"Ingestion to Kusto failed. $ingestionInfoString, " +
+            s"partition: '$partitionId', from: '${ingestionStatusResult.ingestionSourcePath}', " +
+            s"Operation ${ingestionStatusResult.operationId}")
     }
   }
 }

connector/src/main/scala/com/microsoft/kusto/spark/datasink/KustoSinkOptions.scala

@@ -3,12 +3,11 @@
 
 package com.microsoft.kusto.spark.datasink
 
-import java.util.UUID
-import java.util.concurrent.TimeUnit
-
 import com.microsoft.kusto.spark.common.KustoOptions
 import com.microsoft.kusto.spark.utils.KustoConstants
 
+import java.util.UUID
+import java.util.concurrent.TimeUnit
 import scala.concurrent.duration.FiniteDuration
 
 object KustoSinkOptions extends KustoOptions {
@@ -62,6 +61,9 @@ object KustoSinkOptions extends KustoOptions {
   // after which it will move the data to the destination table (the last part is a metadata operation only)
   // If set to 'Queued', the write operation finishes after data is processed by the workers, the data may not be completely
   // available up until the service finishes loading it and failures on the service side will not propagate to Spark.
+  // If set to 'KustoStreaming', Kusto streaming ingestion will be used. Streaming ingestion should be used if latency of less than a few seconds is required
+  // or To optimize operational processing of many tables where the stream of data into each table is relatively small (a few records per second).
+  // If a batch exceeds 4 MB, it will be broken into multiple appropriately sized chunks.
   val KUSTO_WRITE_MODE: String = newOption("writeMode")
 
   // Provide a temporary table name that will be used for this write operation to achieve transactional write and move
@@ -72,6 +74,10 @@ object KustoSinkOptions extends KustoOptions {
   // https://docs.microsoft.com/azure/data-explorer/kusto/management/auto-delete-policy
   // Use this option if you want to persist partial write results (as the failure could be of a single partition)
   val KUSTO_TEMP_TABLE_NAME: String = newOption("tempTableName")
+
+  // The chunk size that we want to use while iterating over "streaming batch". The default is 4MB.
+  // Every streaming ingest will be sent in chunks of this size.
+  val KUSTO_STREAMING_INGEST_SIZE_IN_MB: String = newOption("streamingIngestSizeInMB")
 }
 
 object SinkTableCreationMode extends Enumeration {
@@ -86,28 +92,29 @@ object SchemaAdjustmentMode extends Enumeration {
 
 object WriteMode extends Enumeration {
   type WriteMode = Value
-  val Transactional, Queued = Value
+  val Transactional, Queued, KustoStreaming = Value
 }
 
 case class WriteOptions(
-    pollingOnDriver: Boolean = false,
-    tableCreateOptions: SinkTableCreationMode.SinkTableCreationMode =
+                         pollingOnDriver: Boolean = false,
+                         tableCreateOptions: SinkTableCreationMode.SinkTableCreationMode =
       SinkTableCreationMode.FailIfNotExist,
-    isAsync: Boolean = false,
-    writeResultLimit: String = KustoSinkOptions.NONE_RESULT_LIMIT,
-    timeZone: String = "UTC",
-    timeout: FiniteDuration = new FiniteDuration(
+                         isAsync: Boolean = false,
+                         writeResultLimit: String = KustoSinkOptions.NONE_RESULT_LIMIT,
+                         timeZone: String = "UTC",
+                         timeout: FiniteDuration = new FiniteDuration(
       KustoConstants.DefaultWaitingIntervalLongRunning.toInt,
       TimeUnit.SECONDS),
-    ingestionProperties: Option[String] = None,
-    batchLimit: Int = KustoConstants.DefaultBatchingLimit,
-    requestId: String = UUID.randomUUID().toString,
-    autoCleanupTime: FiniteDuration =
+                         maybeSparkIngestionProperties: Option[SparkIngestionProperties] = None,
+                         batchLimit: Int = KustoConstants.DefaultBatchingLimit,
+                         requestId: String = UUID.randomUUID().toString,
+                         autoCleanupTime: FiniteDuration =
       new FiniteDuration(KustoConstants.DefaultCleaningInterval.toInt, TimeUnit.SECONDS),
-    maxRetriesOnMoveExtents: Int = 10,
-    minimalExtentsCountForSplitMerge: Int = 400,
-    adjustSchema: SchemaAdjustmentMode.SchemaAdjustmentMode = SchemaAdjustmentMode.NoAdjustment,
-    isTransactionalMode: Boolean = true,
-    userTempTableName: Option[String] = None,
-    disableFlushImmediately: Boolean = false,
-    ensureNoDupBlobs: Boolean = false)
+                         maxRetriesOnMoveExtents: Int = 10,
+                         minimalExtentsCountForSplitMerge: Int = 400,
+                         adjustSchema: SchemaAdjustmentMode.SchemaAdjustmentMode = SchemaAdjustmentMode.NoAdjustment,
+                         writeMode: WriteMode.WriteMode = WriteMode.Transactional,
+                         userTempTableName: Option[String] = None,
+                         disableFlushImmediately: Boolean = false,
+                         ensureNoDupBlobs: Boolean = false,
+                         streamIngestUncompressedMaxSize: Int = KustoConstants.DefaultMaxStreamingBytesUncompressed)