Coverage for app/backend/src/tests/sql_linkage.py: 97%

1""" 

2Finds tables in a query that are tied to the rest of it by a value (e.g. `passport_date_of_birth = users.birthdate`) 

3rather than by a key. SQLAlchemy's FROM-linter can't catch this: it counts any predicate as a connection, and never 

4runs on view DDL. 

5""" 

6 

7import itertools 

8import warnings 

9from typing import Any 

10 

11from sqlalchemy import Column, Table 

12from sqlalchemy.sql import operators 

13from sqlalchemy.sql.elements import ( 

14 BinaryExpression, 

15 BooleanClauseList, 

16 ColumnClause, 

17 ColumnElement, 

18 _anonymous_label, 

19) 

20from sqlalchemy.sql.selectable import CompoundSelect, FromClause, Join, Select 

21from sqlalchemy.sql.util import find_tables 

22 

23 

24def _conjuncts(clause: ColumnElement[Any] | None) -> list[ColumnElement[Any]]: 

25 if clause is None: 

26 return [] 

27 if isinstance(clause, BooleanClauseList) and clause.operator is operators.and_: 

28 return [conjunct for sub in clause.clauses for conjunct in _conjuncts(sub)] 

29 return [clause] 

30 

31 

32def _origins(expr: Any) -> set[Column[Any]]: 

33 if not isinstance(expr, ColumnClause): 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 return set() 

35 return {col for col in expr.proxy_set if isinstance(col, Column) and isinstance(col.table, Table)} 

36 

37 

38def _fk_targets(col: Column[Any]) -> set[Column[Any]]: 

39 return {fk.column for fk in col.foreign_keys} 

40 

41 

42def _is_key_link(left: Any, right: Any) -> bool: 

43 left_origins, right_origins = _origins(left), _origins(right) 

44 if left_origins & right_origins: 

45 return True 

46 for a, b in itertools.product(left_origins, right_origins): 

47 if b in _fk_targets(a) or a in _fk_targets(b) or _fk_targets(a) & _fk_targets(b): 

48 return True 

49 return False 

50 

51 

52def _leaves(from_clause: FromClause) -> list[FromClause]: 

53 if isinstance(from_clause, Join): 

54 return _leaves(from_clause.left) + _leaves(from_clause.right) 

55 return [from_clause] 

56 

57 

58def _join_conditions(from_clause: FromClause) -> list[ColumnElement[Any]]: 

59 if isinstance(from_clause, Join): 

60 return [ 

61 *_conjuncts(from_clause.onclause), 

62 *_join_conditions(from_clause.left), 

63 *_join_conditions(from_clause.right), 

64 ] 

65 return [] 

66 

67 

68def _key(from_clause: FromClause) -> str: 

69 # by name rather than identity: the ORM annotates its own copies of tables, so they compare equal but aren't the same 

70 name = getattr(from_clause, "name", None) 

71 return str(name) if name is not None else str(id(from_clause)) 

72 

73 

74def _describe(from_clause: FromClause) -> str: 

75 name = getattr(from_clause, "name", None) 

76 if name is not None and not isinstance(name, _anonymous_label): 

77 return str(name) 

78 inner = getattr(from_clause, "element", None) 

79 tables = sorted({table.name for table in find_tables(inner)}) if inner is not None else [] 

80 return f"unnamed subquery over {', '.join(tables)}" if tables else "unnamed subquery" 

81 

82 

83class _Components: 

84 def __init__(self, keys: list[str]) -> None: 

85 self._parent = {key: key for key in keys} 

86 

87 def find(self, key: str) -> str: 

88 while self._parent[key] != key: 

89 self._parent[key] = self._parent[self._parent[key]] 

90 key = self._parent[key] 

91 return key 

92 

93 def union(self, left: str, right: str) -> None: 

94 self._parent[self.find(left)] = self.find(right) 

95 

96 def count(self) -> int: 

97 return len({self.find(key) for key in self._parent}) 

98 

99 

100def find_unkeyed_joins(statement: Select[Any] | CompoundSelect[Any], path: str = "query") -> list[str]: 

101 """Returns a problem for every FROM list (recursing into subqueries and unions) not connected by keys.""" 

102 if isinstance(statement, CompoundSelect): 

103 return [ 

104 problem 

105 for index, arm in enumerate(statement.selects) 

106 if isinstance(arm, (Select, CompoundSelect)) 

107 for problem in find_unkeyed_joins(arm, f"{path}/union[{index}]") 

108 ] 

109 

110 with warnings.catch_warnings(): 

111 # resolving the FROM list compiles against the default dialect, which warns about DISTINCT ON 

112 warnings.filterwarnings("ignore", "DISTINCT ON is currently supported only by the PostgreSQL dialect") 

113 froms = statement.get_final_froms() 

114 leaves = [leaf for from_clause in froms for leaf in _leaves(from_clause)] 

115 conditions = [ 

116 *(condition for from_clause in froms for condition in _join_conditions(from_clause)), 

117 *_conjuncts(statement.whereclause), 

118 ] 

119 

120 by_key = {_key(leaf): leaf for leaf in leaves} 

121 components = _Components(list(by_key)) 

122 unkeyed: list[BinaryExpression[Any]] = [] 

123 for condition in conditions: 

124 if not isinstance(condition, BinaryExpression) or condition.operator is not operators.eq: 

125 continue 

126 sides = [getattr(getattr(condition, side), "table", None) for side in ("left", "right")] 

127 if any(side is None for side in sides): 

128 continue 

129 left, right = (_key(side) for side in sides) # type: ignore[arg-type] 

130 if left not in by_key or right not in by_key or left == right: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 continue 

132 if _is_key_link(condition.left, condition.right): 

133 components.union(left, right) 

134 else: 

135 unkeyed.append(condition) 

136 

137 problems = [] 

138 if components.count() > 1: 

139 mistaken = "".join(f"\n only linked by: {condition}" for condition in unkeyed) 

140 problems.append( 

141 f"{path}: {', '.join(sorted(_describe(leaf) for leaf in leaves))} are not all connected by keys, so rows " 

142 f"are paired across every combination{mistaken}" 

143 ) 

144 

145 for leaf in leaves: 

146 element = getattr(leaf, "element", None) 

147 if isinstance(element, (Select, CompoundSelect)): 

148 problems += find_unkeyed_joins(element, f"{path}/{_describe(leaf)}") 

149 return problems