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 

28 

29import sqlalchemy as sa 

30from sqlalchemy import orm 

31from sqlalchemy.dialects.postgresql import UUID as PGUUID 

32 

33from wuttjamaican.util import make_true_uuid 

34 

35 

36# nb. this convention comes from upstream docs 

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

38naming_convention = { 

39 'ix': 'ix_%(column_0_label)s', 

40 'uq': 'uq_%(table_name)s_%(column_0_name)s', 

41 'ck': 'ck_%(table_name)s_%(constraint_name)s', 

42 'fk': 'fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s', 

43 'pk': 'pk_%(table_name)s', 

44} 

45 

46 

47class ModelBase: 

48 """ """ 

49 

50 def __iter__(self): 

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

52 state = sa.inspect(self) 

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

54 return iter([(field, getattr(self, field)) 

55 for field in fields]) 

56 

57 def __getitem__(self, key): 

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

59 state = sa.inspect(self) 

60 if hasattr(state.attrs, key): 

61 return getattr(self, key) 

62 

63 

64class UUID(sa.types.TypeDecorator): 

65 """ 

66 Platform-independent UUID type. 

67 

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

69 stringified hex values. 

70 

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

72 documentation 

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

74 """ 

75 impl = sa.CHAR 

76 cache_ok = True 

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

78 

79 def load_dialect_impl(self, dialect): 

80 """ """ 

81 if dialect.name == "postgresql": 

82 return dialect.type_descriptor(PGUUID()) 

83 else: 

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

85 

86 def process_bind_param(self, value, dialect): 

87 """ """ 

88 if value is None: 

89 return value 

90 elif dialect.name == "postgresql": 

91 return str(value) 

92 else: 

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

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

95 else: 

96 # hexstring 

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

98 

99 def process_result_value(self, value, dialect): 

100 """ """ 

101 if value is None: 

102 return value 

103 else: 

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

105 value = _uuid.UUID(value) 

106 return value 

107 

108 

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

110 """ 

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

112 """ 

113 if not args: 

114 args = (UUID(),) 

115 kwargs.setdefault('primary_key', True) 

116 kwargs.setdefault('nullable', False) 

117 kwargs.setdefault('default', make_true_uuid) 

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

119 

120 

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

122 """ 

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

124 

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

126 e.g. ``'user.uuid'``. 

127 """ 

128 if not args: 

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

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

131 

132 

133def make_topo_sortkey(model): 

134 """ 

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

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

137 single class mapper and return a sequence number associated with 

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

139 topological table sorting. 

140 

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

142 containing model classes. 

143 """ 

144 metadata = model.Base.metadata 

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

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

147 

148 def sortkey(name): 

149 cls = getattr(model, name) 

150 mapper = orm.class_mapper(cls) 

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

152 

153 return sortkey