Coverage for python/felis/browser/diagram.py: 97%

1"""The relationship diagram of a schema's foreign keys, drawn by Mermaid in 

2the browser. 

3""" 

4 

5# This file is part of felis. 

6# 

7# Developed for the LSST Data Management System. 

8# This product includes software developed by the LSST Project 

9# (https://www.lsst.org). 

10# See the COPYRIGHT file at the top-level directory of this distribution 

11# for details of code ownership. 

12# 

13# This program is free software: you can redistribute it and/or modify 

14# it under the terms of the GNU General Public License as published by 

15# the Free Software Foundation, either version 3 of the License, or 

16# (at your option) any later version. 

17# 

18# This program is distributed in the hope that it will be useful, 

19# but WITHOUT ANY WARRANTY; without even the implied warranty of 

20# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

21# GNU General Public License for more details. 

22# 

23# You should have received a copy of the GNU General Public License 

24# along with this program. If not, see <https://www.gnu.org/licenses/>. 

25 

26from __future__ import annotations 

27 

28import re 

29 

30from felis.datamodel import ForeignKeyConstraint, Schema, Table 

31 

32from .schemas import _foreign_keys, _primary_key_columns, _resolve_column 

33 

34 

35def _mermaid_label(text: str) -> str: 

36 """Return text for a quoted Mermaid label: " becomes ', and % and 

37 backslashes are removed. 

38 """ 

39 return re.sub(r"[%\\]", "", text.replace('"', "'")) 

40 

41 

42def _mermaid_name(text: str) -> str: 

43 """Return a Mermaid attribute name: letters, digits, _ and -, not starting 

44 with a digit. 

45 """ 

46 name = re.sub(r"[^A-Za-z0-9_-]", "_", text) 

47 if not re.match(r"[A-Za-z_]", name): 

48 name = f"_{name}" 

49 return name 

50 

51 

52class _RelationshipDiagram: 

53 """The Mermaid entity-relationship diagram of a schema's foreign keys. 

54 

55 It shows the tables involved in foreign keys, in schema order, with their 

56 primary and foreign key columns, and one relationship per foreign key, 

57 labelled with its name. Tables are keyed t0, t1, ... in the diagram text. 

58 """ 

59 

60 def __init__(self, schema: Schema) -> None: 

61 self._schema = schema 

62 # Each foreign key, with the table it is defined on. 

63 self._foreign_keys = [ 

64 (table, constraint) for table in schema.tables for constraint in _foreign_keys(table) 

65 ] 

66 

67 def text(self) -> str | None: 

68 """Return the diagram's Mermaid text, or `None` if there are no foreign 

69 keys. 

70 """ 

71 if not self._foreign_keys: 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 return None 

73 tables = self._tables() 

74 keys = {table.name: f"t{number}" for number, table in enumerate(tables)} 

75 lines = ["erDiagram"] 

76 for table in tables: 

77 lines += self._entity_lines(table, keys[table.name]) 

78 for table, constraint in self._foreign_keys: 

79 parent = self._referenced_table(constraint) 

80 label = _mermaid_label(constraint.name) 

81 lines.append(f' {keys[table.name]} }}o--|| {keys[parent.name]} : "{label}"') 

82 return "\n".join(lines) 

83 

84 def _referenced_table(self, constraint: ForeignKeyConstraint) -> Table: 

85 referenced = _resolve_column(self._schema, constraint.referenced_columns[0]) 

86 return self._schema.get_table_by_column(referenced) 

87 

88 def _tables(self) -> list[Table]: 

89 names = set() 

90 for table, constraint in self._foreign_keys: 

91 names.add(table.name) 

92 names.add(self._referenced_table(constraint).name) 

93 return [table for table in self._schema.tables if table.name in names] 

94 

95 def _entity_lines(self, table: Table, key: str) -> list[str]: 

96 primary_key = {column.name for column in _primary_key_columns(self._schema, table)} 

97 foreign_key = set() 

98 for constraint in _foreign_keys(table): 

99 for ref in constraint.columns: 

100 foreign_key.add(_resolve_column(self._schema, ref).name) 

101 lines = [f' {key}["{_mermaid_label(table.name)}"] {{'] 

102 for column in table.columns: 

103 marks = [] 

104 if column.name in primary_key: 

105 marks.append("PK") 

106 if column.name in foreign_key: 

107 marks.append("FK") 

108 if marks: 

109 # Mermaid draws its "type name" pairs in that order; writing 

110 # the name first shows it in the first column. 

111 lines.append(f" {_mermaid_name(column.name)} {column.datatype} {', '.join(marks)}") 

112 lines.append(" }") 

113 return lines