Coverage for /Users/davegaeddert/Development/dropseed/plain/plain-models/plain/models/db.py: 70%

160 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-16 22:04 -0500

1import pkgutil 

2from importlib import import_module 

3 

4from plain import signals 

5from plain.exceptions import ImproperlyConfigured 

6from plain.runtime import settings 

7from plain.utils.connection import BaseConnectionHandler, ConnectionProxy 

8from plain.utils.functional import cached_property 

9from plain.utils.module_loading import import_string 

10 

11DEFAULT_DB_ALIAS = "default" 

12PLAIN_VERSION_PICKLE_KEY = "_plain_version" 

13 

14 

15class Error(Exception): 

16 pass 

17 

18 

19class InterfaceError(Error): 

20 pass 

21 

22 

23class DatabaseError(Error): 

24 pass 

25 

26 

27class DataError(DatabaseError): 

28 pass 

29 

30 

31class OperationalError(DatabaseError): 

32 pass 

33 

34 

35class IntegrityError(DatabaseError): 

36 pass 

37 

38 

39class InternalError(DatabaseError): 

40 pass 

41 

42 

43class ProgrammingError(DatabaseError): 

44 pass 

45 

46 

47class NotSupportedError(DatabaseError): 

48 pass 

49 

50 

51class DatabaseErrorWrapper: 

52 """ 

53 Context manager and decorator that reraises backend-specific database 

54 exceptions using Plain's common wrappers. 

55 """ 

56 

57 def __init__(self, wrapper): 

58 """ 

59 wrapper is a database wrapper. 

60 

61 It must have a Database attribute defining PEP-249 exceptions. 

62 """ 

63 self.wrapper = wrapper 

64 

65 def __enter__(self): 

66 pass 

67 

68 def __exit__(self, exc_type, exc_value, traceback): 

69 if exc_type is None: 

70 return 

71 for plain_exc_type in ( 

72 DataError, 

73 OperationalError, 

74 IntegrityError, 

75 InternalError, 

76 ProgrammingError, 

77 NotSupportedError, 

78 DatabaseError, 

79 InterfaceError, 

80 Error, 

81 ): 

82 db_exc_type = getattr(self.wrapper.Database, plain_exc_type.__name__) 

83 if issubclass(exc_type, db_exc_type): 

84 plain_exc_value = plain_exc_type(*exc_value.args) 

85 # Only set the 'errors_occurred' flag for errors that may make 

86 # the connection unusable. 

87 if plain_exc_type not in (DataError, IntegrityError): 

88 self.wrapper.errors_occurred = True 

89 raise plain_exc_value.with_traceback(traceback) from exc_value 

90 

91 def __call__(self, func): 

92 # Note that we are intentionally not using @wraps here for performance 

93 # reasons. Refs #21109. 

94 def inner(*args, **kwargs): 

95 with self: 

96 return func(*args, **kwargs) 

97 

98 return inner 

99 

100 

101def load_backend(backend_name): 

102 """ 

103 Return a database backend's "base" module given a fully qualified database 

104 backend name, or raise an error if it doesn't exist. 

105 """ 

106 try: 

107 return import_module("%s.base" % backend_name) 

108 except ImportError as e_user: 

109 # The database backend wasn't found. Display a helpful error message 

110 # listing all built-in database backends. 

111 import plain.models.backends 

112 

113 builtin_backends = [ 

114 name 

115 for _, name, ispkg in pkgutil.iter_modules(plain.models.backends.__path__) 

116 if ispkg and name not in {"base", "dummy"} 

117 ] 

118 if backend_name not in [ 

119 "plain.models.backends.%s" % b for b in builtin_backends 

120 ]: 

121 backend_reprs = map(repr, sorted(builtin_backends)) 

122 raise ImproperlyConfigured( 

123 "{!r} isn't an available database backend or couldn't be " 

124 "imported. Check the above exception. To use one of the " 

125 "built-in backends, use 'plain.models.backends.XXX', where XXX " 

126 "is one of:\n" 

127 " {}".format(backend_name, ", ".join(backend_reprs)) 

128 ) from e_user 

129 else: 

130 # If there's some other error, this must be an error in Plain 

131 raise 

132 

133 

134class ConnectionHandler(BaseConnectionHandler): 

135 settings_name = "DATABASES" 

136 

137 def configure_settings(self, databases): 

138 databases = super().configure_settings(databases) 

139 if databases == {}: 

140 databases[DEFAULT_DB_ALIAS] = {"ENGINE": "plain.models.backends.dummy"} 

141 elif DEFAULT_DB_ALIAS not in databases: 

142 raise ImproperlyConfigured( 

143 f"You must define a '{DEFAULT_DB_ALIAS}' database." 

144 ) 

145 elif databases[DEFAULT_DB_ALIAS] == {}: 

146 databases[DEFAULT_DB_ALIAS]["ENGINE"] = "plain.models.backends.dummy" 

147 

148 # Configure default settings. 

149 for conn in databases.values(): 

150 conn.setdefault("AUTOCOMMIT", True) 

151 conn.setdefault("ENGINE", "plain.models.backends.dummy") 

152 if conn["ENGINE"] == "plain.models.backends." or not conn["ENGINE"]: 

153 conn["ENGINE"] = "plain.models.backends.dummy" 

154 conn.setdefault("CONN_MAX_AGE", 0) 

155 conn.setdefault("CONN_HEALTH_CHECKS", False) 

156 conn.setdefault("OPTIONS", {}) 

157 conn.setdefault("TIME_ZONE", None) 

158 for setting in ["NAME", "USER", "PASSWORD", "HOST", "PORT"]: 

159 conn.setdefault(setting, "") 

160 

161 test_settings = conn.setdefault("TEST", {}) 

162 default_test_settings = [ 

163 ("CHARSET", None), 

164 ("COLLATION", None), 

165 ("MIGRATE", True), 

166 ("MIRROR", None), 

167 ("NAME", None), 

168 ] 

169 for key, value in default_test_settings: 

170 test_settings.setdefault(key, value) 

171 return databases 

172 

173 @property 

174 def databases(self): 

175 # Maintained for backward compatibility as some 3rd party packages have 

176 # made use of this private API in the past. It is no longer used within 

177 # Plain itself. 

178 return self.settings 

179 

180 def create_connection(self, alias): 

181 db = self.settings[alias] 

182 backend = load_backend(db["ENGINE"]) 

183 return backend.DatabaseWrapper(db, alias) 

184 

185 

186class ConnectionRouter: 

187 def __init__(self, routers=None): 

188 """ 

189 If routers is not specified, default to settings.DATABASE_ROUTERS. 

190 """ 

191 self._routers = routers 

192 

193 @cached_property 

194 def routers(self): 

195 if self._routers is None: 

196 self._routers = settings.DATABASE_ROUTERS 

197 routers = [] 

198 for r in self._routers: 

199 if isinstance(r, str): 

200 router = import_string(r)() 

201 else: 

202 router = r 

203 routers.append(router) 

204 return routers 

205 

206 def _router_func(action): 

207 def _route_db(self, model, **hints): 

208 chosen_db = None 

209 for router in self.routers: 

210 try: 

211 method = getattr(router, action) 

212 except AttributeError: 

213 # If the router doesn't have a method, skip to the next one. 

214 pass 

215 else: 

216 chosen_db = method(model, **hints) 

217 if chosen_db: 

218 return chosen_db 

219 instance = hints.get("instance") 

220 if instance is not None and instance._state.db: 

221 return instance._state.db 

222 return DEFAULT_DB_ALIAS 

223 

224 return _route_db 

225 

226 db_for_read = _router_func("db_for_read") 

227 db_for_write = _router_func("db_for_write") 

228 

229 def allow_relation(self, obj1, obj2, **hints): 

230 for router in self.routers: 

231 try: 

232 method = router.allow_relation 

233 except AttributeError: 

234 # If the router doesn't have a method, skip to the next one. 

235 pass 

236 else: 

237 allow = method(obj1, obj2, **hints) 

238 if allow is not None: 

239 return allow 

240 return obj1._state.db == obj2._state.db 

241 

242 def allow_migrate(self, db, package_label, **hints): 

243 for router in self.routers: 

244 try: 

245 method = router.allow_migrate 

246 except AttributeError: 

247 # If the router doesn't have a method, skip to the next one. 

248 continue 

249 

250 allow = method(db, package_label, **hints) 

251 

252 if allow is not None: 

253 return allow 

254 return True 

255 

256 def allow_migrate_model(self, db, model): 

257 return self.allow_migrate( 

258 db, 

259 model._meta.package_label, 

260 model_name=model._meta.model_name, 

261 model=model, 

262 ) 

263 

264 def get_migratable_models(self, package_config, db, include_auto_created=False): 

265 """Return app models allowed to be migrated on provided db.""" 

266 models = package_config.get_models(include_auto_created=include_auto_created) 

267 return [model for model in models if self.allow_migrate_model(db, model)] 

268 

269 

270connections = ConnectionHandler() 

271 

272router = ConnectionRouter() 

273 

274# For backwards compatibility. Prefer connections['default'] instead. 

275connection = ConnectionProxy(connections, DEFAULT_DB_ALIAS) 

276 

277 

278# Register an event to reset saved queries when a Plain request is started. 

279def reset_queries(**kwargs): 

280 for conn in connections.all(initialized_only=True): 

281 conn.queries_log.clear() 

282 

283 

284signals.request_started.connect(reset_queries) 

285 

286 

287# Register an event to reset transaction state and close connections past 

288# their lifetime. 

289def close_old_connections(**kwargs): 

290 for conn in connections.all(initialized_only=True): 

291 conn.close_if_unusable_or_obsolete() 

292 

293 

294signals.request_started.connect(close_old_connections) 

295signals.request_finished.connect(close_old_connections)