chore: Scala client for certified events

build.sbt

@@ -36,7 +36,10 @@ val extraDependencies = Seq(
   // Although breeze 1.2 is already provided by Spark, this is needed for Azure Synapse Spark 3.2 pools.
   // Otherwise a NoSuchMethodError will be thrown by interpretability code. This problem only happens
   // to Azure Synapse Spark 3.2 pools.
-  "org.scalanlp" %% "breeze" % "1.2"
+  "org.scalanlp" %% "breeze" % "1.2",
+  "com.typesafe.play" %% "play" % "2.8.8",
+  "com.pauldijou" %% "jwt-core" % "3.0.0",
+  "org.json" % "json" % "20210307"
 ).map(d => d excludeAll (excludes: _*))
 val dependencies = coreDependencies ++ extraDependencies
 

core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/RESTHelpers.scala

@@ -14,6 +14,8 @@ import scala.annotation.tailrec
 import scala.concurrent.blocking
 import scala.util.Try
 
+import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
+
 object RESTHelpers {
   lazy val RequestTimeout = 60000
 
@@ -53,7 +55,6 @@ object RESTHelpers {
                backoffs: List[Int] = List(100, 500, 1000),  //scalastyle:ignore magic.number
                expectedCodes: Set[Int] = Set(),
                close: Boolean = true): CloseableHttpResponse = {
-
     retry(backoffs, { () =>
       val response = Client.execute(request)
       try {
@@ -78,6 +79,10 @@ object RESTHelpers {
         }
       } catch {
         case e: Exception =>
+          println(s"RESTHelpers::safeSend: getting error response parsing." +
+            s". Exception = $e")
+          SynapseMLLogging.logMessage(s"RESTHelpers::safeSend: getting error response parsing." +
+            s". Exception = $e")
           response.close()
           throw e
       } finally {
@@ -89,7 +94,18 @@ object RESTHelpers {
   }
 
   def parseResult(result: CloseableHttpResponse): String = {
-    IOUtils.toString(result.getEntity.getContent, "utf-8")
+    var res: String = ""
+    try {
+      res = IOUtils.toString(result.getEntity.getContent, "utf-8")
+    }
+    catch{
+      case e: Exception =>
+        println(s"RestHelpers::parseResult: getting exception parsing response." +
+          s"Exception = $e")
+        SynapseMLLogging.logMessage(s"RestHelpers::parseResult: getting exception parsing response." +
+          s"Exception = $e")
+    }
+    res
   }
 
   def sendAndParseJson(request: HttpRequestBase, expectedCodes: Set[Int]=Set()): JsValue = {
@@ -98,5 +114,4 @@ object RESTHelpers {
     response.close()
     output
   }
-
 }

core/src/main/scala/com/microsoft/azure/synapse/ml/logging/Usage/FabricTokenServiceClient.scala

@@ -0,0 +1,71 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.logging.Usage
+
+import java.util.UUID
+import java.net.URL
+import java.net.InetAddress
+import java.lang.management.ManagementFactory
+import com.microsoft.azure.synapse.ml.logging.common.WebUtils._
+import spray.json.DefaultJsonProtocol.StringJsonFormat
+import com.microsoft.azure.synapse.ml.logging.Usage.FabricConstants._
+import spray.json.{JsArray, JsObject, JsValue, _}
+import com.microsoft.azure.synapse.ml.logging.common.WebUtils.{usageGet}
+import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
+
+class FabricTokenServiceClient {
+  val resourceMapping = Map(
+    "https://storage.azure.com" -> "storage",
+    "storage" -> "storage",
+    "https://analysis.windows.net/powerbi/api" -> "pbi",
+    "pbi" -> "pbi",
+    "https://vault.azure.net" -> "keyvault",
+    "keyvault" -> "keyvault",
+    "https://kusto.kusto.windows.net" -> "kusto",
+    "kusto" -> "kusto"
+  )
+
+  val hostname = InetAddress.getLocalHost.getHostName
+  val processDetail = ManagementFactory.getRuntimeMXBean().getName()
+  val processName = processDetail.substring(processDetail.indexOf('@') + 1)
+
+  val fabricConbtext = FabricUtils.getFabricContext()
+  val synapseTokenserviceEndpoint = fabricConbtext(SynapseTokenServiceEndpoint)
+  val workloadEndpoint = fabricConbtext(TridentLakehouseTokenServiceEndpoint)
+  val sessionToken = fabricConbtext(TridentSessionToken)
+  val clusterIdentifier = fabricConbtext(SynapseClusterIdentifier)
+
+  def getAccessToken(resourceParam: String): String = {
+    if (!resourceMapping.contains(resourceParam)) {
+      throw new Exception(s"$resourceParam not supported")
+    }
+    val resource = resourceMapping.getOrElse(resourceParam, "")
+    val rid = UUID.randomUUID().toString()
+    //to do workloadEndpoint
+    val targetUrl = new URL(workloadEndpoint)
+    var headers = Map(
+      "x-ms-cluster-identifier" -> clusterIdentifier,
+      "x-ms-workload-resource-moniker" -> clusterIdentifier,
+      "Content-Type" -> "application/json;charset=utf-8",
+      "x-ms-proxy-host" -> s"${targetUrl.getProtocol}://${targetUrl.getHost}",
+      "x-ms-partner-token" -> sessionToken,
+      "User-Agent" -> s"Trident Token Library - HostName:$hostname, ProcessName:$processName",
+      "x-ms-client-request-id" -> rid
+    )
+    var url = s"$synapseTokenserviceEndpoint/api/v1/proxy${targetUrl.getPath}/access?resource=$resource"
+    var response: JsValue = JsonParser("")
+    try {
+      response = usageGet(url, headers)
+      if (response.asJsObject.fields("status_code").convertTo[String] != 200
+        || response.asJsObject.fields("content").convertTo[String].isEmpty) {
+        throw new Exception("Fetch access token error")
+      }
+    } catch {
+      case e: Exception =>
+        println(s"getAccessToken: Failed to fetch cluster details. Exception = $e. (usage test)")
+        SynapseMLLogging.logMessage(s"getAccessToken: Failed to fetch cluster details. Exception = $e. (usage test)")
+    }
+    response.asJsObject.fields("content").toString().getBytes("UTF-8").toString()
+  }
+}

core/src/main/scala/com/microsoft/azure/synapse/ml/logging/Usage/FabricUtils.scala

@@ -0,0 +1,64 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.logging.Usage
+import com.microsoft.azure.synapse.ml.logging.Usage.FabricConstants._
+import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
+import spray.json._
+import spray.json.DefaultJsonProtocol._
+import scala.util.matching.Regex
+import scala.io.Source
+
+case class TokenServiceConfig(tokenServiceEndpoint: String,
+                              clusterType: String,
+                              clusterName: String,
+                              sessionToken: String)
+
+object TokenServiceConfigProtocol extends DefaultJsonProtocol {
+  implicit val TokenServiceConfigFormat: RootJsonFormat[TokenServiceConfig] = jsonFormat4(TokenServiceConfig)
+}
+
+import TokenServiceConfigProtocol._
+
+object FabricUtils {
+  var TridentContext = Map[String, String]()
+
+  def getFabricContext(): Map[String, String] = {
+    if (TridentContext.nonEmpty) {
+      TridentContext
+    } else {
+      try {
+        val lines = scala.io.Source.fromFile(FabricConstants.ContextFilePath).getLines().toList
+        for (line <- lines) {
+          if (line.split('=').length == 2) {
+            val Array(k, v) = line.split('=')
+            TridentContext += (k.trim -> v.trim)
+          }
+        }
+
+        var fileContent: String = Source.fromFile(FabricConstants.TokenServiceFilePath).mkString
+        fileContent = cleanJson(fileContent)
+        val tokenServiceConfigJson = fileContent.parseJson
+
+        // Extract the values from the JSON using Spray JSON's automatic JSON-to-case-class conversion
+        val tokenServiceConfig = tokenServiceConfigJson.convertTo[TokenServiceConfig]
+        // Populate the TridentContext map
+        TridentContext += (FabricConstants.SynapseTokenServiceEndpoint -> tokenServiceConfig.tokenServiceEndpoint)
+        TridentContext += (FabricConstants.SynapseClusterType -> tokenServiceConfig.clusterType)
+        TridentContext += (FabricConstants.SynapseClusterIdentifier -> tokenServiceConfig.clusterName)
+        TridentContext += (FabricConstants.TridentSessionToken -> tokenServiceConfig.sessionToken)
+      } catch {
+        case e: Exception =>
+          SynapseMLLogging.logMessage(s"Error reading Fabric context file: $e")
+          throw e
+      }
+    }
+    TridentContext
+  }
+
+  def cleanJson(s: String): String = {
+    val pattern: Regex = ",[ \t\r\n]+}".r
+    val cleanedJson = pattern.replaceAllIn(s, "}")
+    cleanedJson
+  }
+}

core/src/main/scala/com/microsoft/azure/synapse/ml/logging/Usage/FeatureUsagePayload.scala

@@ -0,0 +1,10 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.logging.Usage
+
+//import scala.collection.mutable.Map
+
+case class FeatureUsagePayload(feature_name: UsageFeatureNames.Value,
+    activity_name: FeatureActivityName.Value,
+    attributes: Map[String, String] = Map.empty[String, String] )

core/src/main/scala/com/microsoft/azure/synapse/ml/logging/Usage/TokenUtils.scala

@@ -0,0 +1,156 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.logging.Usage
+
+import scala.reflect.runtime.currentMirror
+import scala.reflect.runtime.universe._
+import java.time.Instant
+import org.apache.spark.SparkContext
+import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
+import spray.json.DefaultJsonProtocol.{StringJsonFormat, jsonFormat3}
+
+import java.util.UUID
+import com.microsoft.azure.synapse.ml.logging.common.WebUtils._
+
+import java.util.Date
+import pdi.jwt._
+import org.json.JSONObject
+import spray.json.RootJsonFormat
+
+import scala.util.{Failure, Success, Try}
+
+case class MwcToken (TargetUriHost: String, CapacityObjectId: String, Token: String)
+object TokenUtils {
+  var AADToken: String = ""
+  val MwcWorkloadTypeMl = "ML"
+
+  def getAccessToken(): String = {
+    if (checkTokenValid(this.AADToken))
+      this.AADToken
+    else {
+      refreshAccessToken()
+      this.AADToken
+    }
+  }
+
+  def getAccessToken(tokenType: String): String = {
+
+    val objectName = "com.microsoft.azure.trident.tokenlibrary.TokenLibrary"
+    val mirror = currentMirror
+    val module = mirror.staticModule(objectName)
+    val obj = mirror.reflectModule(module).instance
+    val objType = mirror.reflect(obj).symbol.toType
+    val methodName = "getAccessToken"
+    val methodSymbols = objType.decl(TermName(methodName)).asTerm.alternatives
+    val argType = typeOf[String]
+
+    val selectedMethodSymbol = methodSymbols.find { m =>
+      m.asMethod.paramLists match {
+        case List(List(param)) => param.typeSignature =:= argType
+        case _ => false
+      }
+    }.getOrElse(throw new NoSuchMethodException(s"Method $methodName with argument type $argType not found"))
+
+    val methodMirror = mirror.reflect(obj).reflectMethod(selectedMethodSymbol.asMethod)
+    methodMirror(tokenType).asInstanceOf[String]
+  }
+
+  private def checkTokenValid(token: String): Boolean = {
+    try{
+      val expiryDate: Date = getExpiry(token)
+      val expiryEpoch = expiryDate.toInstant.getEpochSecond
+      val now = Instant.now().getEpochSecond
+      now < expiryEpoch - 60
+    }
+    catch
+    {
+      case e: Exception =>
+        SynapseMLLogging.logMessage(s"TokenUtils::checkTokValid: Token {$token} parsing went wrong (usage test). " +
+          s"Exception = $e")
+        false
+    }
+    /*if (token == null || token.isEmpty()) {
+      false
+    }
+    try {
+      val parsedToken = token.parseJson.asJsObject()
+      val expTime = parsedToken.fields("exp").convertTo[Int]
+      val now = Instant.now().getEpochSecond()
+      now < expTime - 60
+    } catch {
+      case e: Exception => {
+        SynapseMLLogging.logMessage(s"TokenUtils::checkTokValid: Token {$token} parsing went wrong (usage test).")
+        false
+      }
+    }*/
+  }
+
+  private def refreshAccessToken(): Unit = {
+    try {
+      if (SparkContext.getOrCreate() != null) {
+        val token = getAccessToken("pbi")
+        AADToken = token
+      } else {
+        val token = new FabricTokenServiceClient().getAccessToken("pbi")
+        AADToken = token
+      }
+    } catch {
+      case e: Exception =>
+        SynapseMLLogging.logMessage(s"refreshAccessTok: failed to refresh pbi tok. Exception: {$e}. (usage test)")
+    }
+  }
+
+  def getMWCToken(shared_host: String, WorkspaceId: String, capacity_id: String,
+                    workload_type: String): MwcToken = {
+    val url: String = shared_host + "/metadata/v201606/generatemwctokenv2"
+
+    val payLoad = s"""{
+      |"capacityObjectId": "$capacity_id",
+      |"workspaceObjectId": "$WorkspaceId",
+      |"workloadType": "$workload_type"
+    }""".stripMargin
+
+    val driverAADToken = getAccessToken()
+
+    val headers = Map(
+      "Content-Type" -> "application/json",
+      "Authorization" -> s"""Bearer $driverAADToken""".stripMargin,
+      "x-ms-workload-resource-moniker" -> UUID.randomUUID().toString
+    )
+
+    try{
+      val response = usagePost(url, payLoad, headers)
+      /*if (response.asJsObject.fields("status_code").convertTo[String] != 200
+        || response.asJsObject.fields("content").convertTo[String].isEmpty) {
+        throw new Exception("Fetch access token error")
+      }*/
+      var targetUriHost = response.asJsObject.fields("TargetUriHost").convertTo[String]
+      targetUriHost = s"https://$targetUriHost"
+      response.asJsObject.fields.updated("TargetUriHost", targetUriHost)
+
+      implicit val mwcTokenFormat: RootJsonFormat[MwcToken] = jsonFormat3(MwcToken)
+      //implicit val mwcTokenFormat = jsonFormat3(MwcToken)
+      response.convertTo[MwcToken]
+    }
+    catch {
+      case e: Exception =>
+        SynapseMLLogging.logMessage(s"getMWCTok: Failed to fetch cluster details: $e. (usage test)")
+        throw e
+    }
+  }
+
+  private def getExpiry(accessToken: String): Date = {
+    val jwtOptions = new JwtOptions(false, false, false, 0)
+    val jwtTokenDecoded: Try[(String, String, String)] = Jwt.decodeRawAll(accessToken, jwtOptions)
+    jwtTokenDecoded match {
+      case Success((_, payload, _)) =>
+        val jsonPayload: JSONObject = new JSONObject(payload)
+        val expiry = jsonPayload.get("exp").toString
+        new Date(expiry.toLong * 1000)
+      case Failure(t) =>
+        SynapseMLLogging.logMessage(t.getMessage)
+        throw t
+    }
+  }
+}