Spark Client StorageID support (#8918)

Author: itaigilo

clients/spark/build.sbt

@@ -27,6 +27,7 @@ Compile / PB.targets := Seq(
 
 testFrameworks += new TestFramework("org.scalameter.ScalaMeterFramework")
 Test / logBuffered := false
+
 // Uncomment to get accurate benchmarks with just "sbt test".
 // Otherwise tell sbt to
 //     "testOnly io.treeverse.clients.ReadSSTableBenchmark"
@@ -42,7 +43,7 @@ buildInfoPackage := "io.treeverse.clients"
 enablePlugins(S3Plugin, BuildInfoPlugin)
 
 libraryDependencies ++= Seq(
-  "io.lakefs" % "sdk" % "1.0.0",
+  "io.lakefs" % "sdk" % "1.53.1",
   "org.apache.spark" %% "spark-sql" % "3.1.2" % "provided",
   "com.thesamet.scalapb" %% "scalapb-runtime" % scalapb.compiler.Version.scalapbVersion % "protobuf",
   "org.apache.hadoop" % "hadoop-aws" % hadoopVersion % "provided",

clients/spark/src/main/scala/io/treeverse/clients/ApiClient.scala

@@ -168,7 +168,7 @@ class ApiClient private (conf: APIConfigurations) {
     val storageNamespace = key.storageClientType match {
       case StorageClientType.HadoopFS =>
         ApiClient
-          .translateURI(URI.create(repo.getStorageNamespace), getBlockstoreType)
+          .translateURI(URI.create(repo.getStorageNamespace), getBlockstoreType(repo.getStorageId))
           .normalize()
           .toString
       case StorageClientType.SDKClient => repo.getStorageNamespace
@@ -210,19 +210,27 @@ class ApiClient private (conf: APIConfigurations) {
     retryWrapper.wrapWithRetry(prepareGcCommits)
   }
 
-  def getGarbageCollectionRules(repoName: String): String = {
-    val getGcRules = new dev.failsafe.function.CheckedSupplier[GarbageCollectionRules]() {
-      def get(): GarbageCollectionRules = repositoriesApi.getGCRules(repoName).execute()
+  def getRepository(repoName: String): Repository = {
+    val getRepo = new dev.failsafe.function.CheckedSupplier[Repository]() {
+      def get(): Repository = repositoriesApi.getRepository(repoName).execute()
     }
-    val gcRules = retryWrapper.wrapWithRetry(getGcRules)
-    gcRules.toString()
+    retryWrapper.wrapWithRetry(getRepo)
   }
 
-  def getBlockstoreType: String = {
+  def getBlockstoreType(storageID: String): String = {
     val getStorageConfig = new dev.failsafe.function.CheckedSupplier[StorageConfig]() {
       def get(): StorageConfig = {
         val cfg = configApi.getConfig.execute()
-        cfg.getStorageConfig
+        val storageConfigList = cfg.getStorageConfigList
+        if (storageConfigList.isEmpty || storageConfigList.size() == 1) {
+          cfg.getStorageConfig
+        } else {
+          storageConfigList.asScala
+            .find(_.getBlockstoreId == storageID)
+            .getOrElse(
+              throw new IllegalArgumentException(s"Storage config not found for ID: $storageID")
+            )
+        }
       }
     }
     val storageConfig = retryWrapper.wrapWithRetry(getStorageConfig)

clients/spark/src/main/scala/io/treeverse/gc/GarbageCollection.scala

@@ -143,7 +143,8 @@ object GarbageCollection {
     val apiConf =
       APIConfigurations(apiURL, accessKey, secretKey, connectionTimeout, readTimeout, sourceName)
     val apiClient = ApiClient.get(apiConf)
-    val storageType = apiClient.getBlockstoreType
+    val storageID = apiClient.getRepository(repo).getStorageId
+    val storageType = apiClient.getBlockstoreType(storageID)
     var storageNamespace = apiClient.getStorageNamespace(repo, StorageClientType.HadoopFS)
     if (!storageNamespace.endsWith("/")) {
       storageNamespace += "/"