dataconnect: gradle plugin: fix DataConnectExecutableDownloadTask for arch-specific executable versions

firebase-dataconnect/gradleplugin/plugin/src/main/kotlin/com/google/firebase/dataconnect/gradle/plugin/DataConnectExecutableDownloadTask.kt

@@ -15,6 +15,8 @@
  */
 package com.google.firebase.dataconnect.gradle.plugin
 
+import io.github.z4kn4fein.semver.Version
+import io.github.z4kn4fein.semver.toVersion
 import java.io.File
 import java.net.HttpURLConnection
 import java.net.URL
@@ -46,6 +48,8 @@ abstract class DataConnectExecutableDownloadTask : DefaultTask() {
 
   @get:Input @get:Optional abstract val operatingSystem: Property<OperatingSystem>
 
+  @get:Input @get:Optional abstract val cpuArchitecture: Property<CpuArchitecture>
+
   @get:Internal abstract val buildDirectory: DirectoryProperty
 
   @get:OutputFile abstract val outputFile: RegularFileProperty
@@ -59,12 +63,14 @@ abstract class DataConnectExecutableDownloadTask : DefaultTask() {
     val inputFile: File? = inputFile.orNull?.asFile
     val version: String? = version.orNull
     val operatingSystem: OperatingSystem = operatingSystem.get()
+    val cpuArchitecture: CpuArchitecture = cpuArchitecture.get()
     val buildDirectory: File = buildDirectory.get().asFile
     val outputFile: File = outputFile.get().asFile
 
     logger.info("inputFile: {}", inputFile)
     logger.info("version: {}", version)
     logger.info("operatingSystem: {}", operatingSystem)
+    logger.info("cpuArchitecture: {}", cpuArchitecture)
     logger.info("buildDirectory: {}", buildDirectory)
     logger.info("outputFile: {}", outputFile)
 
@@ -81,8 +87,15 @@ abstract class DataConnectExecutableDownloadTask : DefaultTask() {
     } else if (inputFile !== null) {
       runWithFile(inputFile = inputFile, outputFile = outputFile)
     } else if (version !== null) {
-      downloadDataConnectExecutable(version, operatingSystem, outputFile, execOperations)
-      verifyOutputFile(outputFile, operatingSystem, version)
+      val parsedVersion = version.toVersion()
+      downloadDataConnectExecutable(
+        parsedVersion,
+        operatingSystem,
+        cpuArchitecture,
+        outputFile,
+        execOperations,
+      )
+      verifyOutputFile(outputFile, operatingSystem, cpuArchitecture, parsedVersion)
     } else {
       throw DataConnectGradleException(
         "chc94cq7vx",
@@ -95,35 +108,34 @@ abstract class DataConnectExecutableDownloadTask : DefaultTask() {
   private fun verifyOutputFile(
     outputFile: File,
     operatingSystem: OperatingSystem,
-    version: String
+    cpuArchitecture: CpuArchitecture,
+    version: Version
   ) {
     logger.info("Verifying file size and SHA512 digest of file: {}", outputFile)
     val fileInfo = FileInfo.forFile(outputFile)
 
-    val allVersions = DataConnectExecutableVersionsRegistry.load().versions
-    val allVersionNames =
-      allVersions
-        .asSequence()
-        .filter { it.os == operatingSystem }
-        .map { it.version }
-        .distinct()
-        .sorted()
-        .joinToString(", ")
-    val applicableVersions =
-      allVersions.filter { it.version.toString() == version && it.os == operatingSystem }
+    val versionsForOsArch: List<DataConnectExecutableVersionsRegistry.VersionInfo> =
+      DataConnectExecutableVersionsRegistry.load().versions.filter {
+        it.os == operatingSystem && (it.arch === null || it.arch == cpuArchitecture)
+      }
+    val applicableVersions = versionsForOsArch.filter { it.version == version }
+    val versionNamesStringForOsArch =
+      versionsForOsArch.map { it.version }.distinct().sorted().joinToString(", ") { it.toString() }
 
+    val osDisplayId: String =
+      operatingSystem.downloadFileNameMoniker + "/" + cpuArchitecture.downloadFileNameMoniker
     if (applicableVersions.isEmpty()) {
       val message =
         "verification information for Data Connect toolkit executable" +
-          " version $version for $operatingSystem is not known;" +
-          " known versions for $operatingSystem are: $allVersionNames" +
+          " version $version for $osDisplayId is not known;" +
+          " known versions for $osDisplayId are: $versionNamesStringForOsArch" +
           " (loaded from ${DataConnectExecutableVersionsRegistry.PATH})"
       logger.error("ERROR: $message")
       throw DataConnectGradleException("ym8assbfgw", message)
     } else if (applicableVersions.size > 1) {
       val message =
         "INTERNAL ERROR: ${applicableVersions.size} verification information records for" +
-          " Data Connect toolkit executable version $version for $operatingSystem were found in" +
+          " Data Connect toolkit executable version $version for $osDisplayId were found in" +
           " ${DataConnectExecutableVersionsRegistry.PATH}, but expected exactly 1"
       logger.error("ERROR: $message")
       throw DataConnectGradleException("zyw5xrky6e", message)
@@ -152,8 +164,8 @@ abstract class DataConnectExecutableDownloadTask : DefaultTask() {
     if (verificationErrors.isNotEmpty()) {
       val errorMessage =
         "Verification of ${outputFile.absolutePath}" +
-          " (version=${versionInfo.version} os=${versionInfo.os}) failed:" +
-          " ${verificationErrors.joinToString(", ")}"
+          " (version=${versionInfo.version} os=${versionInfo.os} arch=${versionInfo.arch})" +
+          " failed: ${verificationErrors.joinToString(", ")}"
       logger.error(errorMessage)
       throw DataConnectGradleException("x9dfwhjr9c", errorMessage)
     }
@@ -199,19 +211,42 @@ abstract class DataConnectExecutableDownloadTask : DefaultTask() {
   }
 
   companion object {
+
+    private val firstVersionIncludingArchInFileName = "2.16.0".toVersion()
+
+    private val OperatingSystem.downloadFileNameMoniker: String
+      get() =
+        when (this) {
+          OperatingSystem.Windows -> "windows"
+          OperatingSystem.MacOS -> "macos"
+          OperatingSystem.Linux -> "linux"
+        }
+
+    private val CpuArchitecture.downloadFileNameMoniker: String
+      get() =
+        when (this) {
+          CpuArchitecture.AMD64 -> "amd64"
+          CpuArchitecture.ARM64 -> "arm64"
+        }
+
     fun Task.downloadDataConnectExecutable(
-      version: String,
+      version: Version,
       operatingSystem: OperatingSystem,
+      cpuArchitecture: CpuArchitecture,
       outputFile: File,
       execOperations: ExecOperations
     ) {
-      val osName =
-        when (operatingSystem) {
-          OperatingSystem.Windows -> "windows"
-          OperatingSystem.MacOS -> "macos"
-          OperatingSystem.Linux -> "linux"
+      val downloadFileName = buildString {
+        append("dataconnect-emulator")
+        append('-')
+        append(operatingSystem.downloadFileNameMoniker)
+        if (version >= firstVersionIncludingArchInFileName) {
+          append('-')
+          append(cpuArchitecture.downloadFileNameMoniker)
         }
-      val downloadFileName = "dataconnect-emulator-$osName-v$version"
+        append("-v")
+        append(version)
+      }
       val url =
         URL("https://storage.googleapis.com/firemat-preview-drop/emulator/$downloadFileName")
 

firebase-dataconnect/gradleplugin/plugin/src/main/kotlin/com/google/firebase/dataconnect/gradle/plugin/DataConnectGradlePlugin.kt

@@ -111,6 +111,7 @@ abstract class DataConnectGradlePlugin : Plugin<Project> {
             )
           )
           operatingSystem.set(dataConnectProviders.operatingSystem)
+          cpuArchitecture.set(dataConnectProviders.cpuArchitecture)
           outputFile.set(
             dataConnectExecutable.map {
               when (it) {

firebase-dataconnect/gradleplugin/plugin/src/main/kotlin/com/google/firebase/dataconnect/gradle/plugin/DataConnectProviders.kt

@@ -62,6 +62,8 @@ class DataConnectProviders(
 
   val operatingSystem: Provider<OperatingSystem> = project.provider { OperatingSystem.current() }
 
+  val cpuArchitecture: Provider<CpuArchitecture> = project.provider { CpuArchitecture.current() }
+
   val postgresConnectionUrl: Provider<String> = run {
     val gradlePropertyName = "dataconnect.emulator.postgresConnectionUrl"
     val valueFromLocalSettings: Provider<String> = localSettings.postgresConnectionUrl