Spawn missing tables on cascading RestrGraph (#1368)

* Spawn missing tables on cascading RestrGraph

* Update changelog

* Check alias node child in `_is_out` util

* PR feedback

Author: CBroz1

CHANGELOG.md

@@ -23,6 +23,10 @@ import all foreign key references.
 - Allow email send on space check success, clean up maintenance logging #1381
 - Update pynwb pin to >=2.5.0 for `TimeSeries.get_timestamps` #1385
 
+### Infrastructure
+
+- Auto-load within-Spyglass tables for graph operations #1368
+
 ### Pipelines
 
 - Behavior

src/spyglass/utils/dj_graph.py

@@ -12,7 +12,7 @@
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Set, Tuple, Union
 
-from datajoint import FreeTable, Table
+from datajoint import FreeTable, Table, VirtualModule
 from datajoint import config as dj_config
 from datajoint.condition import make_condition
 from datajoint.hash import key_hash
@@ -283,12 +283,48 @@ def _get_ft(self, table, with_restr=False, warn=True):
 
         return ft & restr
 
+    def _has_out_prefix(self, table):
+        return (
+            table.split(".")[0].split("_")[0].strip("`") not in SHARED_MODULES
+        )
+
+    def _spawn_virtual_module(self, table):
+        schema = table.split(".")[0].strip("`")
+        logger.warning(f"Spawning tables for {schema}")
+        vm = VirtualModule(f"RestrGraph_{schema}", schema)
+        v_graph = vm.schema.connection.dependencies
+        v_graph.load()
+
+        self.graph.add_nodes_from(v_graph.nodes(data=True))
+        self.graph.add_edges_from(v_graph.edges(data=True))
+
     def _is_out(self, table, warn=True):
         """Check if table is outside of spyglass."""
         table = ensure_names(table)
-        if table in self.graph.nodes:
+        if table.isnumeric():  # if alias node, determine status from child
+            children = list(self.graph.children(table))
+            if len(children) > 1:
+                raise ValueError(f"Alias has multiple connections: {table}")
+            if children[0].isnumeric():
+                raise ValueError(f"Alias of alias, should not happen: {table}")
+            return self._is_out(children[0])
+
+        # If already in imported, return
+        # Reverts #1356: was `table in self.graph.nodes`, now `get`
+        #   - Present nodes may be children of imported, with no data
+        #   - Only imported tables have data retrieved by `get`
+        if self.graph.nodes.get(table):
             return False
-        ret = table.split(".")[0].split("_")[0].strip("`") not in SHARED_MODULES
+
+        # If within spyglass, attempt spawn
+        ret = self._has_out_prefix(table)
+        if not ret:
+            _ = self._spawn_virtual_module(table)
+
+        # If spawn successful, return
+        if self.graph.nodes.get(table):
+            return False
+
         if warn and ret:  # Log warning if outside
             logger.warning(f"Skipping unimported: {table}")  # pragma: no cover
         return ret

src/spyglass/utils/dj_mixin.py

@@ -853,6 +853,29 @@ def restrict_by(
 
         return ret
 
+    def restrict_all(
+        self,
+        restriction: str = True,
+        direction: str = "down",
+        return_graph: bool = False,
+        verbose: bool = False,
+    ) -> List:
+        RestrGraph = self._graph_deps[1]
+        rg = RestrGraph(
+            seed_table=self,
+            leaves=dict(
+                table_name=self.full_table_name,
+                restriction=self.restriction or restriction,
+            ),
+            direction=direction,
+            banned_tables=list(self._banned_search_tables),
+            cascade=True,
+            verbose=verbose,
+        )
+        if return_graph:
+            return rg
+        logger.info(rg.restr_ft)
+
     # ------------------------------ Check locks ------------------------------
 
     def exec_sql_fetchall(self, query):