Allow BXLs to aquery actions directly; used for action debugging.

facebook#1250

Author: josephsumabat

app/buck2_build_api/src/query/bxl.rs

@@ -12,6 +12,7 @@ use std::future::Future;
 use std::pin::Pin;
 
 use async_trait::async_trait;
+use buck2_artifact::actions::key::ActionKey;
 use buck2_core::cells::CellResolver;
 use buck2_core::cells::name::CellName;
 use buck2_core::configuration::compatibility::MaybeCompatible;
@@ -182,6 +183,11 @@ pub trait BxlAqueryFunctions: Send {
         dice: &mut DiceComputations<'_>,
         targets: &TargetSet<ActionQueryNode>,
     ) -> buck2_error::Result<TargetSet<ActionQueryNode>>;
+    async fn get_action_nodes(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        action_keys: Vec<ActionKey>,
+    ) -> buck2_error::Result<TargetSet<ActionQueryNode>>;
 }
 
 pub static NEW_BXL_CQUERY_FUNCTIONS: LateBinding<

app/buck2_bxl/src/bxl/starlark_defs/aquery.rs

@@ -9,6 +9,7 @@
  */
 
 use allocative::Allocative;
+use buck2_artifact::actions::key::ActionKey;
 use buck2_build_api::actions::query::ActionQueryNode;
 use buck2_build_api::query::bxl::BxlAqueryFunctions;
 use buck2_build_api::query::bxl::NEW_BXL_AQUERY_FUNCTIONS;
@@ -46,6 +47,7 @@ use starlark::values::type_repr::StarlarkTypeRepr;
 
 use crate::bxl::starlark_defs::context::BxlContext;
 use crate::bxl::starlark_defs::context::ErrorPrinter;
+use crate::bxl::starlark_defs::nodes::action::StarlarkAction;
 use crate::bxl::starlark_defs::nodes::action::StarlarkActionQueryNode;
 use crate::bxl::starlark_defs::providers_expr::AnyProvidersExprArg;
 use crate::bxl::starlark_defs::providers_expr::ProvidersExpr;
@@ -123,6 +125,7 @@ enum UnpackActionNodes<'v> {
     ActionQueryNodesSet(&'v StarlarkTargetSet<ActionQueryNode>),
     ConfiguredProviders(AnyProvidersExprArg<'v>),
     ConfiguredTargets(ConfiguredTargetListExprArg<'v>),
+    StarlarkActions(UnpackList<StarlarkAction>),
 }
 
 // Aquery operates on `ActionQueryNode`s. Under the hood, the target set of action query nodes is obtained
@@ -141,6 +144,14 @@ async fn unpack_action_nodes<'v>(
             return Ok(action_nodes.into_iter().map(|v| v.0).collect());
         }
         UnpackActionNodes::ActionQueryNodesSet(action_nodes) => return Ok(action_nodes.0.clone()),
+        UnpackActionNodes::StarlarkActions(actions) => {
+            let action_keys: Vec<ActionKey> = actions
+                .into_iter()
+                .map(|starlark_action| starlark_action.0.key().dupe())
+                .collect();
+
+            return aquery_env.get_action_nodes(dice, action_keys).await;
+        }
         UnpackActionNodes::ConfiguredProviders(arg) => {
             ProvidersExpr::<ConfiguredProvidersLabel>::unpack(
                 arg,
@@ -183,6 +194,11 @@ async fn unpack_action_nodes<'v>(
 ///
 /// Query results are `target_set`s of `action_query_node`s, which supports iteration,
 /// indexing, `len()`, set addition/subtraction, and `equals()`.
+///
+/// Actions can be specified as:
+/// - Target expressions (configured targets/providers)
+/// - Existing action query nodes or target sets
+/// - `bxl.Action` objects (obtained from `ctx.audit().output()`)
 #[starlark_module]
 fn aquery_methods(builder: &mut MethodsBuilder) {
     /// The deps query for finding the transitive closure of dependencies.

app/buck2_query_impls/src/aquery/bxl.rs

@@ -12,6 +12,7 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use buck2_artifact::actions::key::ActionKey;
 use buck2_build_api::actions::query::ActionQueryNode;
 use buck2_build_api::analysis::calculation::RuleAnalysisCalculation;
 use buck2_build_api::query::bxl::BxlAqueryFunctions;
@@ -302,6 +303,24 @@ impl BxlAqueryFunctions for BxlAqueryFunctionsImpl {
         })
         .await
     }
+
+    async fn get_action_nodes(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        action_keys: Vec<ActionKey>,
+    ) -> buck2_error::Result<TargetSet<ActionQueryNode>> {
+        dice.with_linear_recompute(|dice| async move {
+            let delegate = self.aquery_delegate(&dice).await?;
+            let mut result = TargetSet::new();
+            let nodes = buck2_util::future::try_join_all(
+                action_keys.iter().map(|key| delegate.get_action_node(&key)),
+            )
+            .await?;
+            result.extend(nodes);
+            Ok(result)
+        })
+        .await
+    }
 }
 
 pub(crate) fn init_new_bxl_aquery_functions() {

tests/core/query/aquery/test_aquery.py

@@ -154,3 +154,16 @@ async def test_bxl_aquery_eval(buck: Buck) -> None:
 @buck_test()
 async def test_bxl_aquery_action_query_node(buck: Buck) -> None:
     await buck.bxl("//:aquery.bxl:action_query_node")
+
+
+# Tests for bxl.Action support in aquery operations
+@buck_test()
+async def test_bxl_action_deps_0(buck: Buck) -> None:
+    """Test passing a bxl.Action to aquery.deps() with depth=0 returns itself"""
+    await buck.bxl("//:aquery.bxl:action_deps_0")
+
+
+@buck_test()
+async def test_bxl_action_deps(buck: Buck) -> None:
+    """Test that you can isolate the deps of one action by itself"""
+    await buck.bxl("//:aquery.bxl:action_deps")

tests/core/query/aquery/test_aquery_data/aquery.bxl

@@ -141,3 +141,38 @@ action_query_node = bxl_main(
     impl = _impl_action_query_node,
     cli_args = {},
 )
+
+def _impl_action_deps_0(ctx):
+    """Test passing a bxl.Action to aquery.deps() returns itself"""
+    nodes = ctx.aquery().eval("//:test")
+
+    action = nodes[0].action()
+
+    # Pass the bxl.Action rather than a target to aquery.deps()
+    result = ctx.aquery().deps([action], depth = 0)
+    _assert_eq(action.outputs(), result[0].action().outputs())
+
+action_deps_0 = bxl_main(
+    impl = _impl_action_deps_0,
+    cli_args = {},
+)
+
+def _impl_action_deps(ctx):
+    """Test that you can isolate the deps of one action by itself"""
+    nodes = ctx.aquery().eval("//:test")
+
+    action = nodes[0].action()
+
+    # Pass the bxl.Action rather than a target to aquery.deps()
+    result = ctx.aquery().deps([action], depth = 1)
+
+    _assert_eq(len(result), 2)
+    # First should be itself
+    _assert_eq(action.outputs(), result[0].action().outputs())
+    dep = result[1]
+    _assert_eq(dep.action().outputs()[0].short_path, "dep")
+
+action_deps = bxl_main(
+    impl = _impl_action_deps,
+    cli_args = {},
+)

tests/core/query/aquery/test_aquery_data/prelude.bzl

@@ -9,6 +9,7 @@
 def _test(ctx: AnalysisContext):
     dep = ctx.actions.write("dep", "")
     default = ctx.actions.copy_file("default", dep)
+    other_default = ctx.actions.copy_file("other_default", dep)
     other = ctx.actions.write("other", "")
 
     sub_default = ctx.actions.write("sub_default", "")
@@ -18,7 +19,7 @@ def _test(ctx: AnalysisContext):
     ctx.actions.write("unused", "")
 
     return [DefaultInfo(
-        default_outputs = [default],
+        default_outputs = [default, other_default],
         other_outputs = [other],
         sub_targets = {
             "sub": [