BazelQueryService.kt

package com.bazel_diff.bazel

import com.bazel_diff.log.Logger
import com.bazel_diff.process.Redirect
import com.bazel_diff.process.process
import com.google.devtools.build.lib.analysis.AnalysisProtosV2
import com.google.devtools.build.lib.query2.proto.proto2api.Build
import java.io.File
import java.nio.file.Files
import java.nio.file.Path
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.runBlocking
import org.koin.core.component.KoinComponent
import org.koin.core.component.inject

private val versionComparator =
    compareBy<Triple<Int, Int, Int>> { it.first }.thenBy { it.second }.thenBy { it.third }

class BazelQueryService(
    private val workingDirectory: Path,
    private val bazelPath: Path,
    private val startupOptions: List<String>,
    private val commandOptions: List<String>,
    private val cqueryOptions: List<String>,
    private val keepGoing: Boolean,
    private val noBazelrc: Boolean,
) : KoinComponent {
  private val logger: Logger by inject()
  private val version: Triple<Int, Int, Int> by lazy { runBlocking { determineBazelVersion() } }

  @OptIn(ExperimentalCoroutinesApi::class)
  private suspend fun determineBazelVersion(): Triple<Int, Int, Int> {
    val cmd = arrayOf(bazelPath.toString(), "version")
    logger.i { "Executing Bazel version command: ${cmd.joinToString()}" }
    val result =
        process(
            *cmd,
            stdout = Redirect.CAPTURE,
            workingDirectory = workingDirectory.toFile(),
            stderr = Redirect.PRINT,
            destroyForcibly = true,
        )

    if (result.resultCode != 0) {
      throw RuntimeException("Bazel version command failed, exit code ${result.resultCode}")
    }

    // "bazel version" outputs "Build label: X.Y.Z" on one of the lines; accept that or legacy "bazel X.Y.Z".
    val versionString =
        result.output
            .firstOrNull { it.startsWith("Build label: ") }
            ?.removePrefix("Build label: ")?.trim()
            ?: result.output
                .firstOrNull { it.startsWith("bazel ") }
                ?.removePrefix("bazel ")?.trim()
            ?: throw RuntimeException(
                "Bazel version command returned unexpected output: ${result.output}")
    // Trim off any prerelease suffixes (e.g., 8.6.0-rc1 or 8.6.0rc1).
    val version =
        versionString.split('-')[0].split('.').map { it.takeWhile { c -> c.isDigit() }.toInt() }.toTypedArray()
    return Triple(version[0], version[1], version[2])
  }

  // Use streamed_proto output for cquery if available. This is more efficient than the proto
  // output.
  // https://github.com/bazelbuild/bazel/commit/607d0f7335f95aa0ee236ba3c18ce2a232370cdb
  private val canUseStreamedProtoWithCquery
    get() = versionComparator.compare(version, Triple(7, 0, 0)) >= 0

  // Use an output file for (c)query if supported. This avoids excessively large stdout, which is
  // sent out on the BES.
  // https://github.com/bazelbuild/bazel/commit/514e9052f2c603c53126fbd9436bdd3ad3a1b0c7
  private val canUseOutputFile
    get() = versionComparator.compare(version, Triple(8, 2, 0)) >= 0

  // Bazel 8.6.0+ / 9.0.1+ supports `bazel mod show_repo --output=streamed_proto`, which
  // outputs Build.Repository protos for bzlmod-managed external repos.
  // https://github.com/bazelbuild/bazel/pull/28010
  val canUseBzlmodShowRepo
    get() =
        versionComparator.compare(version, Triple(8, 6, 0)) >= 0 &&
            // 9.0.0 does not have the feature; it landed in 9.0.1.
            version != Triple(9, 0, 0)

  suspend fun query(query: String, useCquery: Boolean = false): List<BazelTarget> {
    // Unfortunately, there is still no direct way to tell if a target is compatible or not with the
    // proto output
    // by itself. So we do an extra cquery with the trick at
    // https://bazel.build/extending/platforms#cquery-incompatible-target-detection to first find
    // all compatible
    // targets.
    val compatibleTargetSet =
        if (useCquery) {
          runQuery(query, useCquery = true, outputCompatibleTargets = true).useLines {
            it.filter { it.isNotBlank() }.toSet()
          }
        } else {
          emptySet()
        }
    val outputFile = runQuery(query, useCquery)

    val targets =
        outputFile.inputStream().buffered().use { proto ->
          if (useCquery) {
            if (canUseStreamedProtoWithCquery) {
                  mutableListOf<AnalysisProtosV2.CqueryResult>()
                      .apply {
                        while (true) {
                          val result =
                              AnalysisProtosV2.CqueryResult.parseDelimitedFrom(proto) ?: break
                          // EOF
                          add(result)
                        }
                      }
                      .flatMap { it.resultsList }
                } else {
                  AnalysisProtosV2.CqueryResult.parseFrom(proto).resultsList
                }
                .mapNotNull { toBazelTarget(it.target) }
                .filter { it.name in compatibleTargetSet }
          } else {
            mutableListOf<Build.Target>()
                .apply {
                  while (true) {
                    val target = Build.Target.parseDelimitedFrom(proto) ?: break
                    // EOF
                    add(target)
                  }
                }
                .mapNotNull { toBazelTarget(it) }
          }
        }

    return targets
  }

  @OptIn(ExperimentalCoroutinesApi::class)
  private suspend fun runQuery(
      query: String,
      useCquery: Boolean,
      outputCompatibleTargets: Boolean = false
  ): File {
    val queryFile = Files.createTempFile(null, ".txt").toFile()
    queryFile.deleteOnExit()
    val outputFile = Files.createTempFile(null, ".bin").toFile()
    outputFile.deleteOnExit()

    queryFile.writeText(query)

    val allowedExitCodes = mutableListOf(0)

    val cmd: MutableList<String> =
        ArrayList<String>().apply {
          add(bazelPath.toString())
          if (noBazelrc) {
            add("--bazelrc=/dev/null")
          }
          addAll(startupOptions)
          if (useCquery) {
            add("cquery")
            if (!outputCompatibleTargets) {
              // There is no need to query the transitions when querying for compatible targets.
              add("--transitions=lite")
            }
          } else {
            add("query")
          }
          add("--output")
          if (useCquery) {
            if (outputCompatibleTargets) {
              add("starlark")
              add("--starlark:file")
              val cqueryStarlarkFile = Files.createTempFile(null, ".cquery").toFile()
              cqueryStarlarkFile.deleteOnExit()
              cqueryStarlarkFile.writeText(
                  """
                    def format(target):
                        if providers(target) == None:
                            return ""
                        if "IncompatiblePlatformProvider" not in providers(target):
                            target_repr = repr(target)
                            if "<alias target" in target_repr:
                                return target_repr.split(" ")[2]
                            return str(target.label)
                        return ""
                    """
                      .trimIndent())
              add(cqueryStarlarkFile.toString())
            } else {
              add(if (canUseStreamedProtoWithCquery) "streamed_proto" else "proto")
            }
          } else {
            add("streamed_proto")
          }
          if (!useCquery) {
            add("--order_output=no")
          }
          if (keepGoing) {
            add("--keep_going")
            allowedExitCodes.add(3)
          }
          if (useCquery) {
            addAll(cqueryOptions)
            add("--consistent_labels")
          } else {
            addAll(commandOptions)
          }
          add("--query_file")
          add(queryFile.toString())
          if (canUseOutputFile) {
            add("--output_file")
            add(outputFile.toString())
          }
        }

    logger.i { "Executing Query: $query" }
    logger.i { "Command: ${cmd.toTypedArray().joinToString()}" }
    val result =
        process(
            *cmd.toTypedArray(),
            stdout = if (canUseOutputFile) Redirect.SILENT else Redirect.ToFile(outputFile),
            workingDirectory = workingDirectory.toFile(),
            stderr = Redirect.PRINT,
            destroyForcibly = true,
        )

    if (!allowedExitCodes.contains(result.resultCode)) {
      logger.w { "Bazel query failed, output: ${result.output.joinToString("\n")}" }
      throw RuntimeException(
          "Bazel query failed, exit code ${result.resultCode}, allowed exit codes: ${allowedExitCodes.joinToString()}")
    }
    return outputFile
  }

  /**
   * Queries bzlmod-managed external repo definitions using `bazel mod show_repo`.
   * Requires Bazel 8.6.0+ or 9.0.1+ which supports `--output=streamed_proto` for this command.
   *
   * The approach:
   * 1. Run `bazel mod dump_repo_mapping ""` to discover the root module's apparent→canonical
   *    repo name mapping (e.g., "bazel_diff_maven" → "rules_jvm_external++maven+maven").
   * 2. Run `bazel mod show_repo @@<canonical>... --output=streamed_proto` to get Repository
   *    proto definitions for each repo (works for both module repos and extension-generated repos).
   * 3. Create synthetic `//external:<apparent_name>` targets for each repo. This matches how
   *    `transformRuleInput` in BazelRule.kt collapses `@apparent_name//...` deps to
   *    `//external:apparent_name`, so the hashing pipeline can detect changes.
   */
  @OptIn(ExperimentalCoroutinesApi::class)
  suspend fun queryBzlmodRepos(): List<BazelTarget> {
    check(canUseBzlmodShowRepo) { "queryBzlmodRepos requires Bazel 8.6.0+ or 9.0.1+" }

    // Step 1: Get the root module's apparent → canonical repo mapping.
    val repoMapping = discoverRepoMapping()
    if (repoMapping.isEmpty()) {
      logger.w { "No repo mappings discovered, skipping mod show_repo" }
      return emptyList()
    }
    logger.i { "Discovered ${repoMapping.size} repo mappings" }

    // Build reverse map: canonical → list of apparent names
    val canonicalToApparent = mutableMapOf<String, MutableList<String>>()
    for ((apparent, canonical) in repoMapping) {
      canonicalToApparent.getOrPut(canonical) { mutableListOf() }.add(apparent)
    }

    // Step 2: Fetch repo definitions via `mod show_repo @@<canonical>... --output=streamed_proto`.
    val canonicalNames = canonicalToApparent.keys.map { "@@$it" }
    val outputFile = Files.createTempFile(null, ".bin").toFile()
    outputFile.deleteOnExit()

    val cmd: MutableList<String> =
        ArrayList<String>().apply {
          add(bazelPath.toString())
          if (noBazelrc) {
            add("--bazelrc=/dev/null")
          }
          addAll(startupOptions)
          add("mod")
          add("show_repo")
          addAll(canonicalNames)
          add("--output=streamed_proto")
        }

    logger.i { "Querying bzlmod repos: ${cmd.joinToString()}" }
    val result =
        process(
            *cmd.toTypedArray(),
            stdout = Redirect.ToFile(outputFile),
            workingDirectory = workingDirectory.toFile(),
            stderr = Redirect.PRINT,
            destroyForcibly = true,
        )

    if (result.resultCode != 0) {
      logger.w { "bazel mod show_repo failed (exit code ${result.resultCode}), skipping bzlmod repos" }
      return emptyList()
    }

    // Step 3: Parse Build.Repository messages and create synthetic targets for each apparent name.
    val repos =
        outputFile.inputStream().buffered().use { proto ->
          mutableListOf<Build.Repository>().apply {
            while (true) {
              val repo = Build.Repository.parseDelimitedFrom(proto) ?: break
              add(repo)
            }
          }
        }

    val targets = mutableListOf<BazelTarget.Rule>()
    for (repo in repos) {
      val apparentNames = canonicalToApparent[repo.canonicalName]
      if (apparentNames != null) {
        for (apparentName in apparentNames) {
          targets.add(repositoryToTarget(repo, apparentName))
        }
      } else {
        // Fallback: use canonical name if no apparent name mapping exists
        targets.add(repositoryToTarget(repo, repo.canonicalName))
      }
    }

    logger.i { "Parsed ${repos.size} bzlmod repos → ${targets.size} synthetic targets" }
    return targets
  }

  /**
   * Converts a Build.Repository proto into a synthetic BazelTarget.Rule named
   * `//external:<targetName>`. This mirrors how WORKSPACE repos appear as `//external:*`
   * targets, and matches the names produced by `transformRuleInput` in BazelRule.kt.
   */
  private fun repositoryToTarget(repo: Build.Repository, targetName: String): BazelTarget.Rule {
    val ruleClass = repo.repoRuleName.ifEmpty { "bzlmod_repo" }

    val target =
        Build.Target.newBuilder()
            .setType(Build.Target.Discriminator.RULE)
            .setRule(
                Build.Rule.newBuilder()
                    .setName("//external:$targetName")
                    .setRuleClass(ruleClass)
                    .addAllAttribute(repo.attributeList))
            .build()
    return BazelTarget.Rule(target)
  }

  /**
   * Discovers the root module's apparent→canonical repo name mapping by running
   * `bazel mod dump_repo_mapping ""`. Returns a map of apparent name → canonical name.
   * Filters out internal repos (bazel_tools, _builtins, local_config_*) that aren't
   * relevant for dependency hashing.
   */
  @OptIn(ExperimentalCoroutinesApi::class)
  private suspend fun discoverRepoMapping(): Map<String, String> {
    val cmd: MutableList<String> =
        ArrayList<String>().apply {
          add(bazelPath.toString())
          if (noBazelrc) {
            add("--bazelrc=/dev/null")
          }
          addAll(startupOptions)
          add("mod")
          add("dump_repo_mapping")
          // Empty string = root module's repo mapping
          add("")
        }

    logger.i { "Discovering repo mapping: ${cmd.joinToString()}" }
    val result =
        process(
            *cmd.toTypedArray(),
            stdout = Redirect.CAPTURE,
            workingDirectory = workingDirectory.toFile(),
            stderr = Redirect.PRINT,
            destroyForcibly = true,
        )

    if (result.resultCode != 0) {
      logger.w { "bazel mod dump_repo_mapping failed (exit code ${result.resultCode})" }
      return emptyMap()
    }

    return try {
      val mapping = mutableMapOf<String, String>()
      for (line in result.output) {
        val trimmed = line.trim()
        if (trimmed.isEmpty()) continue
        val json = com.google.gson.JsonParser.parseString(trimmed).asJsonObject
        for ((apparent, canonicalElem) in json.entrySet()) {
          val canonical = canonicalElem.asString
          // Skip internal/infrastructure repos not relevant for dependency hashing.
          if (apparent.isEmpty() ||
              canonical.isEmpty() ||
              canonical.startsWith("bazel_tools") ||
              canonical.startsWith("_builtins") ||
              canonical.startsWith("local_config_") ||
              canonical.startsWith("rules_java_builtin") ||
              apparent == "bazel_tools" ||
              apparent == "local_config_platform")
              continue
          mapping[apparent] = canonical
        }
      }
      mapping
    } catch (e: Exception) {
      logger.w { "Failed to parse dump_repo_mapping output: ${e.message}" }
      emptyMap()
    }
  }

  private fun toBazelTarget(target: Build.Target): BazelTarget? {
    return when (target.type) {
      Build.Target.Discriminator.RULE -> BazelTarget.Rule(target)
      Build.Target.Discriminator.SOURCE_FILE -> BazelTarget.SourceFile(target)
      Build.Target.Discriminator.GENERATED_FILE -> BazelTarget.GeneratedFile(target)
      else -> {
        logger.w { "Unsupported target type in the build graph: ${target.type.name}" }
        null
      }
    }
  }
}