Coverage for .tox/coverage/lib/python3.11/site-packages/wuttjamaican/db/util.py: 100%

1# -*- coding: utf-8; -*- 

2################################################################################ 

3# 

4# WuttJamaican -- Base package for Wutta Framework 

5# Copyright © 2023-2025 Lance Edgar 

6# 

7# This file is part of Wutta Framework. 

8# 

9# Wutta Framework is free software: you can redistribute it and/or modify it 

10# under the terms of the GNU General Public License as published by the Free 

11# Software Foundation, either version 3 of the License, or (at your option) any 

12# later version. 

13# 

14# Wutta Framework is distributed in the hope that it will be useful, but 

15# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 

16# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for 

17# more details. 

18# 

19# You should have received a copy of the GNU General Public License along with 

20# Wutta Framework. If not, see <http://www.gnu.org/licenses/>. 

21# 

22################################################################################ 

23""" 

24Database Utilities 

25""" 

26 

27import uuid as _uuid 

28from importlib.metadata import version 

29 

30from packaging.version import Version 

31import sqlalchemy as sa 

32from sqlalchemy import orm 

33from sqlalchemy.dialects.postgresql import UUID as PGUUID 

34 

35from wuttjamaican.util import make_true_uuid 

36 

37 

38# nb. this convention comes from upstream docs 

39# https://docs.sqlalchemy.org/en/14/core/constraints.html#constraint-naming-conventions 

40naming_convention = { 

41 "ix": "ix_%(column_0_label)s", 

42 "uq": "uq_%(table_name)s_%(column_0_name)s", 

43 "ck": "ck_%(table_name)s_%(constraint_name)s", 

44 "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", 

45 "pk": "pk_%(table_name)s", 

46} 

47 

48 

49SA2 = True # pylint: disable=invalid-name 

50if Version(version("SQLAlchemy")) < Version("2"): # pragma: no cover 

51 SA2 = False # pylint: disable=invalid-name 

52 

53 

54class ModelBase: # pylint: disable=empty-docstring 

55 """ """ 

56 

57 def __iter__(self): 

58 # nb. we override this to allow for `dict(self)` 

59 state = sa.inspect(self) 

60 fields = [attr.key for attr in state.attrs] 

61 return iter([(field, getattr(self, field)) for field in fields]) 

62 

63 def __getitem__(self, key): 

64 # nb. we override this to allow for `x = self['field']` 

65 state = sa.inspect(self) 

66 if hasattr(state.attrs, key): 

67 return getattr(self, key) 

68 raise KeyError( 

69 f"{self.__class__.__name__} instance has no attr with key: {key}" 

70 ) 

71 

72 

73class UUID( 

74 sa.types.TypeDecorator 

75): # pylint: disable=abstract-method,too-many-ancestors 

76 """ 

77 Platform-independent UUID type. 

78 

79 Uses PostgreSQL's UUID type, otherwise uses CHAR(32), storing as 

80 stringified hex values. 

81 

82 This type definition is based on example from the `SQLAlchemy 

83 documentation 

84 <https://docs.sqlalchemy.org/en/14/core/custom_types.html#backend-agnostic-guid-type>`_. 

85 """ 

86 

87 impl = sa.CHAR 

88 cache_ok = True 

89 """ """ # nb. suppress sphinx autodoc for cache_ok 

90 

91 def load_dialect_impl(self, dialect): # pylint: disable=empty-docstring 

92 """ """ 

93 if dialect.name == "postgresql": 

94 return dialect.type_descriptor(PGUUID()) 

95 return dialect.type_descriptor(sa.CHAR(32)) 

96 

97 def process_bind_param(self, value, dialect): # pylint: disable=empty-docstring 

98 """ """ 

99 if value is None: 

100 return value 

101 

102 if dialect.name == "postgresql": 

103 return str(value) 

104 

105 if not isinstance(value, _uuid.UUID): 

106 value = _uuid.UUID(value) 

107 

108 # hexstring 

109 return f"{value.int:032x}" 

110 

111 def process_result_value( 

112 self, value, dialect 

113 ): # pylint: disable=unused-argument,empty-docstring 

114 """ """ 

115 if value is None: 

116 return value 

117 if not isinstance(value, _uuid.UUID): 

118 value = _uuid.UUID(value) 

119 return value 

120 

121 

122def uuid_column(*args, **kwargs): 

123 """ 

124 Returns a UUID column for use as a table's primary key. 

125 """ 

126 if not args: 

127 args = (UUID(),) 

128 kwargs.setdefault("primary_key", True) 

129 kwargs.setdefault("nullable", False) 

130 kwargs.setdefault("default", make_true_uuid) 

131 if kwargs["primary_key"]: 

132 kwargs.setdefault("doc", "UUID primary key for the table.") 

133 return sa.Column(*args, **kwargs) 

134 

135 

136def uuid_fk_column(target_column, *args, **kwargs): 

137 """ 

138 Returns a UUID column for use as a foreign key to another table. 

139 

140 :param target_column: Name of the table column on the remote side, 

141 e.g. ``'user.uuid'``. 

142 """ 

143 if not args: 

144 args = (UUID(), sa.ForeignKey(target_column)) 

145 return sa.Column(*args, **kwargs) 

146 

147 

148def make_topo_sortkey(model): 

149 """ 

150 Returns a function suitable for use as a ``key`` kwarg to a 

151 standard Python sorting call. This key function will expect a 

152 single class mapper and return a sequence number associated with 

153 that model. The sequence is determined by SQLAlchemy's 

154 topological table sorting. 

155 

156 :param model: Usually the :term:`app model`, but can be any module 

157 containing model classes. 

158 """ 

159 metadata = model.Base.metadata 

160 tables = {table.name: i for i, table in enumerate(metadata.sorted_tables, 1)} 

161 

162 def sortkey(name): 

163 cls = getattr(model, name) 

164 mapper = orm.class_mapper(cls) 

165 return tuple(tables[t.name] for t in mapper.tables) 

166 

167 return sortkey