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-2024 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 

29from packaging.version import Version 

30 

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 

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

51 SA2 = False 

52 

53 

54class ModelBase: 

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)) 

62 for field in fields]) 

63 

64 def __getitem__(self, key): 

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

66 state = sa.inspect(self) 

67 if hasattr(state.attrs, key): 

68 return getattr(self, key) 

69 

70 

71class UUID(sa.types.TypeDecorator): 

72 """ 

73 Platform-independent UUID type. 

74 

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

76 stringified hex values. 

77 

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

79 documentation 

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

81 """ 

82 impl = sa.CHAR 

83 cache_ok = True 

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

85 

86 def load_dialect_impl(self, dialect): 

87 """ """ 

88 if dialect.name == "postgresql": 

89 return dialect.type_descriptor(PGUUID()) 

90 else: 

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

92 

93 def process_bind_param(self, value, dialect): 

94 """ """ 

95 if value is None: 

96 return value 

97 elif dialect.name == "postgresql": 

98 return str(value) 

99 else: 

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

101 return "%.32x" % _uuid.UUID(value).int 

102 else: 

103 # hexstring 

104 return "%.32x" % value.int 

105 

106 def process_result_value(self, value, dialect): 

107 """ """ 

108 if value is None: 

109 return value 

110 else: 

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

112 value = _uuid.UUID(value) 

113 return value 

114 

115 

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

117 """ 

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

119 """ 

120 if not args: 

121 args = (UUID(),) 

122 kwargs.setdefault('primary_key', True) 

123 kwargs.setdefault('nullable', False) 

124 kwargs.setdefault('default', make_true_uuid) 

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

126 

127 

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

129 """ 

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

131 

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

133 e.g. ``'user.uuid'``. 

134 """ 

135 if not args: 

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

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

138 

139 

140def make_topo_sortkey(model): 

141 """ 

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

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

144 single class mapper and return a sequence number associated with 

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

146 topological table sorting. 

147 

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

149 containing model classes. 

150 """ 

151 metadata = model.Base.metadata 

152 tables = dict([(table.name, i) 

153 for i, table in enumerate(metadata.sorted_tables, 1)]) 

154 

155 def sortkey(name): 

156 cls = getattr(model, name) 

157 mapper = orm.class_mapper(cls) 

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

159 

160 return sortkey