Coverage for graphqler / chains / chain_generator.py: 92%

1"""ChainGenerator: orchestrates chain generation using a pluggable strategy.""" 

2 

3from pathlib import Path 

4 

5import networkx 

6import yaml 

7 

8from graphqler import config 

9from graphqler.chains.chain import Chain 

10from graphqler.chains.strategies.base_strategy import BaseChainStrategy 

11from graphqler.chains.strategies.all_dependencies_strategy import AllDependenciesChainStrategy 

12from graphqler.graph.node import Node 

13 

14# All possible mutation_type values assigned to Mutation nodes 

15_ALL_MUTATION_TYPES = ["CREATE", "UPDATE", "DELETE", "UNKNOWN"] 

16 

17 

18class ChainGenerator: 

19 """Generates and stores dependency chains for later inspection and fuzzer consumption. 

20 

21 The generator applies a **3-pass** strategy (mirroring the original DFS passes) when 

22 ``config.DISABLE_MUTATIONS`` is ``False`` (the default): 

23 

24 * **Pass 1** — chains containing only CREATE / QUERY nodes (UPDATE, DELETE, UNKNOWN filtered out) 

25 * **Pass 2** — chains allowing CREATE, QUERY, and UPDATE (DELETE, UNKNOWN filtered out) 

26 * **Pass 3** — all chains (no filter) 

27 

28 All three passes are concatenated into a single list so that the fuzzer can simply 

29 iterate through ``chains`` without any additional filtering. 

30 

31 When ``config.DISABLE_MUTATIONS`` is ``True``, only Query (and Object) chains are 

32 produced — all mutation nodes are excluded entirely. 

33 

34 Usage:: 

35 

36 generator = ChainGenerator() 

37 chains = generator.generate(dependency_graph, starter_nodes) 

38 # chains are also accessible afterwards: 

39 print(generator.chains) 

40 

41 The default strategy is :class:`AllDependenciesChainStrategy`. Pass a different 

42 :class:`BaseChainStrategy` subclass to use an alternative generation method 

43 (e.g. :class:`~graphqler.chains.strategies.dfs_strategy.DFSChainStrategy`). 

44 """ 

45 

46 def __init__(self, strategy: BaseChainStrategy | None = None): 

47 """Initialise the generator with an optional strategy. 

48 

49 Args: 

50 strategy (BaseChainStrategy | None): Chain generation strategy. 

51 Defaults to :class:`DFSChainStrategy` when *None*. 

52 """ 

53 self._strategy: BaseChainStrategy = strategy if strategy is not None else AllDependenciesChainStrategy() 

54 self._chains: list[Chain] = [] 

55 

56 @property 

57 def chains(self) -> list[Chain]: 

58 """The chains produced by the most recent :meth:`generate` call (empty until then).""" 

59 return self._chains 

60 

61 def generate(self, graph: networkx.DiGraph, starter_nodes: list[Node]) -> list[Chain]: 

62 """Generate chains and store them for later inspection. 

63 

64 Applies the 3-pass filtering strategy internally so that the fuzzer receives a 

65 single ordered list of chains ready to execute sequentially. 

66 

67 Args: 

68 graph (networkx.DiGraph): The compiled dependency graph. 

69 starter_nodes (list[Node]): Root nodes to start generation from. 

70 

71 Returns: 

72 list[Chain]: The generated chains (same object as :attr:`chains`). 

73 """ 

74 if config.DISABLE_MUTATIONS: 

75 # Only produce chains that contain Query/Object nodes — exclude all mutations 

76 self._chains = self._strategy.generate( 

77 graph, starter_nodes, filter_mutation_type=_ALL_MUTATION_TYPES 

78 ) 

79 else: 

80 # Pass 1: CREATE + QUERY only (filter UPDATE, DELETE, UNKNOWN) 

81 pass1 = self._strategy.generate( 

82 graph, starter_nodes, 

83 filter_mutation_type=["UPDATE", "DELETE", "UNKNOWN"], 

84 ) 

85 # Pass 2: CREATE + QUERY + UPDATE (filter DELETE, UNKNOWN) 

86 pass2 = self._strategy.generate( 

87 graph, starter_nodes, 

88 filter_mutation_type=["DELETE", "UNKNOWN"], 

89 ) 

90 # Pass 3: all nodes 

91 pass3 = self._strategy.generate( 

92 graph, starter_nodes, 

93 filter_mutation_type=[], 

94 ) 

95 self._chains = pass1 + pass2 + pass3 

96 

97 return self._chains 

98 

99 def save_to_yaml(self, save_path: str) -> None: 

100 """Persist the generated chains to a YAML file for human inspection and optional editing. 

101 

102 Each chain is stored as a list of node names. On reload the names are 

103 resolved back to :class:`~graphqler.graph.node.Node` objects using the 

104 dependency graph. 

105 

106 Args: 

107 save_path (str): Root output directory (same directory used for compilation). 

108 """ 

109 chains_path = Path(save_path) / config.CHAINS_FILE_NAME 

110 chains_path.parent.mkdir(parents=True, exist_ok=True) 

111 data = [{"nodes": [n.name for n in chain.nodes]} for chain in self._chains] 

112 with open(chains_path, "w") as f: 

113 yaml.dump(data, f, default_flow_style=False, sort_keys=False) 

114 

115 def load_from_yaml(self, save_path: str, graph: networkx.DiGraph) -> list[Chain]: 

116 """Load chains from a previously saved YAML file and populate :attr:`chains`. 

117 

118 Node names in the YAML are resolved to :class:`~graphqler.graph.node.Node` 

119 objects using *graph*. Any name that no longer exists in the graph is 

120 silently skipped so that hand-edited files do not crash the fuzzer. 

121 

122 Args: 

123 save_path (str): Root output directory (same directory used for compilation). 

124 graph (networkx.DiGraph): The dependency graph used to look up Node objects. 

125 

126 Returns: 

127 list[Chain]: The loaded chains (same object as :attr:`chains`). 

128 """ 

129 chains_path = Path(save_path) / config.CHAINS_FILE_NAME 

130 if not chains_path.exists(): 

131 self._chains = [] 

132 return self._chains 

133 

134 node_map: dict[str, Node] = {node.name: node for node in graph.nodes()} 

135 with open(chains_path, "r") as f: 

136 data = yaml.safe_load(f) or [] 

137 

138 self._chains = [ 

139 Chain(nodes=[node_map[name] for name in entry.get("nodes", []) if name in node_map]) 

140 for entry in data 

141 ] 

142 return self._chains 

143