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 # nb. this does *not* include association proxy values; 

60 # see also wuttjamaican.util.get_value() 

61 state = sa.inspect(self) 

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

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

64 

65 def __getitem__(self, key): 

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

67 state = sa.inspect(self) 

68 if hasattr(state.attrs, key): 

69 return getattr(self, key) 

70 raise KeyError( 

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

72 ) 

73 

74 

75class UUID( 

76 sa.types.TypeDecorator 

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

78 """ 

79 Platform-independent UUID type. 

80 

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

82 stringified hex values. 

83 

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

85 documentation 

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

87 """ 

88 

89 impl = sa.CHAR 

90 cache_ok = True 

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

92 

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

94 """ """ 

95 if dialect.name == "postgresql": 

96 return dialect.type_descriptor(PGUUID()) 

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

98 

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

100 """ """ 

101 if value is None: 

102 return value 

103 

104 if dialect.name == "postgresql": 

105 return str(value) 

106 

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

108 value = _uuid.UUID(value) 

109 

110 # hexstring 

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

112 

113 def process_result_value( 

114 self, value, dialect 

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

116 """ """ 

117 if value is None: 

118 return value 

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

120 value = _uuid.UUID(value) 

121 return value 

122 

123 

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

125 """ 

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

127 """ 

128 if not args: 

129 args = (UUID(),) 

130 kwargs.setdefault("primary_key", True) 

131 kwargs.setdefault("nullable", False) 

132 kwargs.setdefault("default", make_true_uuid) 

133 if kwargs["primary_key"]: 

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

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

136 

137 

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

139 """ 

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

141 

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

143 e.g. ``'user.uuid'``. 

144 """ 

145 if not args: 

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

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

148 

149 

150def make_topo_sortkey(model): 

151 """ 

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

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

154 single class mapper and return a sequence number associated with 

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

156 topological table sorting. 

157 

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

159 containing model classes. 

160 """ 

161 metadata = model.Base.metadata 

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

163 

164 def sortkey(name): 

165 cls = getattr(model, name) 

166 mapper = orm.class_mapper(cls) 

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

168 

169 return sortkey