spack find: add options for local/upstream only (#42999)

Users requested an option to filter between local/upstream results in `spack find` output.

```
# default behavior, same as without --install-tree argument
$ spack find --install-tree all

# show only local results
$ spack find --install-tree local  

# show results from all upstreams
$ spack find --install-tree upstream 

# show results from a particular upstream or the local install_tree
$ spack find --install-tree /path/to/install/tree/root
```

---------

Co-authored-by: becker33 <becker33@users.noreply.github.com>

6 files changed, 62 insertions, 9 deletions

diff --git a/lib/spack/spack/cmd/find.py b/lib/spack/spack/cmd/find.py
index 070ac9bd0e..f8e74f262f 100644
--- a/lib/spack/spack/cmd/find.py
+++ b/lib/spack/spack/cmd/find.py
@@ -140,6 +140,12 @@ def setup_parser(subparser):
     subparser.add_argument(
         "--only-deprecated", action="store_true", help="show only deprecated packages"
     )
+    subparser.add_argument(
+        "--install-tree",
+        action="store",
+        default="all",
+        help="Install trees to query: 'all' (default), 'local', 'upstream', upstream name or path",
+    )
 
     subparser.add_argument("--start-date", help="earliest date of installation [YYYY-MM-DD]")
     subparser.add_argument("--end-date", help="latest date of installation [YYYY-MM-DD]")
@@ -168,6 +174,12 @@ def query_arguments(args):
 
     q_args = {"installed": installed, "known": known, "explicit": explicit}
 
+    install_tree = args.install_tree
+    upstreams = spack.config.get("upstreams", {})
+    if install_tree in upstreams.keys():
+        install_tree = upstreams[install_tree]["install_tree"]
+    q_args["install_tree"] = install_tree
+
     # Time window of installation
     for attribute in ("start_date", "end_date"):
         date = getattr(args, attribute)
diff --git a/lib/spack/spack/database.py b/lib/spack/spack/database.py
index c2f0567d22..e2f9e57e82 100644
--- a/lib/spack/spack/database.py
+++ b/lib/spack/spack/database.py
@@ -1621,15 +1621,32 @@ class Database:
     query_local.__doc__ += _QUERY_DOCSTRING
 
     def query(self, *args, **kwargs):
-        """Query the Spack database including all upstream databases."""
+        """Query the Spack database including all upstream databases.
+
+        Additional Arguments:
+            install_tree (str): query 'all' (default), 'local', 'upstream', or upstream path
+        """
+        install_tree = kwargs.pop("install_tree", "all")
+        valid_trees = ["all", "upstream", "local", self.root] + [u.root for u in self.upstream_dbs]
+        if install_tree not in valid_trees:
+            msg = "Invalid install_tree argument to Database.query()\n"
+            msg += f"Try one of {', '.join(valid_trees)}"
+            tty.error(msg)
+            return []
+
         upstream_results = []
-        for upstream_db in self.upstream_dbs:
+        upstreams = self.upstream_dbs
+        if install_tree not in ("all", "upstream"):
+            upstreams = [u for u in self.upstream_dbs if u.root == install_tree]
+        for upstream_db in upstreams:
             # queries for upstream DBs need to *not* lock - we may not
             # have permissions to do this and the upstream DBs won't know about
             # us anyway (so e.g. they should never uninstall specs)
             upstream_results.extend(upstream_db._query(*args, **kwargs) or [])
 
-        local_results = set(self.query_local(*args, **kwargs))
+        local_results = []
+        if install_tree in ("all", "local") or self.root == install_tree:
+            local_results = set(self.query_local(*args, **kwargs))
 
         results = list(local_results) + list(x for x in upstream_results if x not in local_results)
 
diff --git a/lib/spack/spack/test/cmd/find.py b/lib/spack/spack/test/cmd/find.py
index 5d551e6df9..e42a938e38 100644
--- a/lib/spack/spack/test/cmd/find.py
+++ b/lib/spack/spack/test/cmd/find.py
@@ -64,6 +64,7 @@ def test_query_arguments():
         implicit=False,
         start_date="2018-02-23",
         end_date=None,
+        install_tree="all",
     )
 
     q_args = query_arguments(args)
@@ -75,6 +76,7 @@ def test_query_arguments():
     assert q_args["explicit"] is any
     assert "start_date" in q_args
     assert "end_date" not in q_args
+    assert q_args["install_tree"] == "all"
 
     # Check that explicit works correctly
     args.explicit = True
diff --git a/lib/spack/spack/test/database.py b/lib/spack/spack/test/database.py
index e02f549532..0b611e93e3 100644
--- a/lib/spack/spack/test/database.py
+++ b/lib/spack/spack/test/database.py
@@ -56,11 +56,31 @@ def upstream_and_downstream_db(tmpdir, gen_mock_layout):
     yield upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout
 
 
+@pytest.mark.parametrize(
+    "install_tree,result",
+    [("all", ["b", "c"]), ("upstream", ["c"]), ("local", ["b"]), ("{u}", ["c"]), ("{d}", ["b"])],
+)
+def test_query_by_install_tree(
+    install_tree, result, upstream_and_downstream_db, mock_packages, monkeypatch, config
+):
+    up_write_db, up_db, up_layout, down_db, down_layout = upstream_and_downstream_db
+
+    # Set the upstream DB to contain "c" and downstream to contain "b")
+    b = spack.spec.Spec("b").concretized()
+    c = spack.spec.Spec("c").concretized()
+    up_write_db.add(c, up_layout)
+    up_db._read()
+    down_db.add(b, down_layout)
+
+    specs = down_db.query(install_tree=install_tree.format(u=up_db.root, d=down_db.root))
+    assert [s.name for s in specs] == result
+
+
 def test_spec_installed_upstream(
     upstream_and_downstream_db, mock_custom_repository, config, monkeypatch
 ):
     """Test whether Spec.installed_upstream() works."""
-    (upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout) = (
+    upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout = (
         upstream_and_downstream_db
     )
 
@@ -86,7 +106,7 @@ def test_spec_installed_upstream(
 
 @pytest.mark.usefixtures("config")
 def test_installed_upstream(upstream_and_downstream_db, tmpdir):
-    (upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout) = (
+    upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout = (
         upstream_and_downstream_db
     )
 
@@ -124,7 +144,7 @@ def test_installed_upstream(upstream_and_downstream_db, tmpdir):
 
 @pytest.mark.usefixtures("config")
 def test_removed_upstream_dep(upstream_and_downstream_db, tmpdir):
-    (upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout) = (
+    upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout = (
         upstream_and_downstream_db
     )
 
@@ -156,7 +176,7 @@ def test_add_to_upstream_after_downstream(upstream_and_downstream_db, tmpdir):
     DB. When a package is recorded as installed in both, the results should
     refer to the downstream DB.
     """
-    (upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout) = (
+    upstream_write_db, upstream_db, upstream_layout, downstream_db, downstream_layout = (
         upstream_and_downstream_db
     )
 
diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash
index 1319d1685f..dc0d1987a6 100755
--- a/share/spack/spack-completion.bash
+++ b/share/spack/spack-completion.bash
@@ -1179,7 +1179,7 @@ _spack_fetch() {
 _spack_find() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help --format -H --hashes --json -d --deps -p --paths --groups --no-groups -l --long -L --very-long -t --tag -N --namespaces -c --show-concretized -f --show-flags --show-full-compiler -x --explicit -X --implicit -u --unknown -m --missing -v --variants --loaded -M --only-missing --deprecated --only-deprecated --start-date --end-date"
+        SPACK_COMPREPLY="-h --help --format -H --hashes --json -d --deps -p --paths --groups --no-groups -l --long -L --very-long -t --tag -N --namespaces -c --show-concretized -f --show-flags --show-full-compiler -x --explicit -X --implicit -u --unknown -m --missing -v --variants --loaded -M --only-missing --deprecated --only-deprecated --install-tree --start-date --end-date"
     else
         _installed_packages
     fi
diff --git a/share/spack/spack-completion.fish b/share/spack/spack-completion.fish
index 47e6fccb31..ff7c3b3ba1 100755
--- a/share/spack/spack-completion.fish
+++ b/share/spack/spack-completion.fish
@@ -1721,7 +1721,7 @@ complete -c spack -n '__fish_spack_using_command fetch' -l deprecated -f -a conf
 complete -c spack -n '__fish_spack_using_command fetch' -l deprecated -d 'allow concretizer to select deprecated versions'
 
 # spack find
-set -g __fish_spack_optspecs_spack_find h/help format= H/hashes json d/deps p/paths groups no-groups l/long L/very-long t/tag= N/namespaces c/show-concretized f/show-flags show-full-compiler x/explicit X/implicit u/unknown m/missing v/variants loaded M/only-missing deprecated only-deprecated start-date= end-date=
+set -g __fish_spack_optspecs_spack_find h/help format= H/hashes json d/deps p/paths groups no-groups l/long L/very-long t/tag= N/namespaces c/show-concretized f/show-flags show-full-compiler x/explicit X/implicit u/unknown m/missing v/variants loaded M/only-missing deprecated only-deprecated install-tree= start-date= end-date=
 complete -c spack -n '__fish_spack_using_command_pos_remainder 0 find' -f -a '(__fish_spack_installed_specs)'
 complete -c spack -n '__fish_spack_using_command find' -s h -l help -f -a help
 complete -c spack -n '__fish_spack_using_command find' -s h -l help -d 'show this help message and exit'
@@ -1771,6 +1771,8 @@ complete -c spack -n '__fish_spack_using_command find' -l deprecated -f -a depre
 complete -c spack -n '__fish_spack_using_command find' -l deprecated -d 'show deprecated packages as well as installed specs'
 complete -c spack -n '__fish_spack_using_command find' -l only-deprecated -f -a only_deprecated
 complete -c spack -n '__fish_spack_using_command find' -l only-deprecated -d 'show only deprecated packages'
+complete -c spack -n '__fish_spack_using_command find' -l install-tree -r -f -a install_tree
+complete -c spack -n '__fish_spack_using_command find' -l install-tree -r -d 'Install trees to query: \'all\' (default), \'local\', \'upstream\', upstream name or path'
 complete -c spack -n '__fish_spack_using_command find' -l start-date -r -f -a start_date
 complete -c spack -n '__fish_spack_using_command find' -l start-date -r -d 'earliest date of installation [YYYY-MM-DD]'
 complete -c spack -n '__fish_spack_using_command find' -l end-date -r -f -a end_date