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

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

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

4# WuttJamaican -- Base package for Wutta Framework

7# This file is part of Wutta Framework.

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

24WuttJamaican - database configuration

25"""

27from collections import OrderedDict

29import sqlalchemy as sa

30from alembic.config import Config as AlembicConfig

31from alembic.script import ScriptDirectory

32from alembic.migration import MigrationContext

34from wuttjamaican.util import load_object, parse_bool, parse_list

37def get_engines(config, prefix):

38 """

39 Construct and return all database engines defined for a given

40 config prefix.

42 For instance if you have a config file with:

44 .. code-block:: ini

46 [wutta.db]

47 keys = default, host

48 default.url = sqlite:///tmp/default.sqlite

49 host.url = sqlite:///tmp/host.sqlite

51 And then you call this function to get those DB engines::

53 get_engines(config, 'wutta.db')

55 The result of that will be like::

57 {'default': Engine(bind='sqlite:///tmp/default.sqlite'),

58 'host': Engine(bind='sqlite:///tmp/host.sqlite')}

60 :param config: App config object.

62 :param prefix: Prefix for the config "section" which contains DB

63 connection info.

65 :returns: A dictionary of SQLAlchemy engines, with keys matching

66 those found in config.

67 """

68 keys = config.get(f"{prefix}.keys", usedb=False)

69 if keys:

70 keys = parse_list(keys)

71 else:

72 keys = ["default"]

74 make_engine = config.get_engine_maker()

76 engines = OrderedDict()

77 cfg = config.get_dict(prefix)

78 for key in keys:

79 key = key.strip()

80 try:

81 engines[key] = make_engine(cfg, prefix=f"{key}.")

82 except KeyError:

83 if key == "default":

84 try:

85 engines[key] = make_engine(cfg, prefix="sqlalchemy.")

86 except KeyError:

87 pass

88 return engines

91def get_setting(session, name):

92 """

93 Get a setting value from the DB.

95 Note that this assumes (for now?) the DB contains a table named

96 ``setting`` with ``(name, value)`` columns.

98 :param session: App DB session.

100 :param name: Name of the setting to get.

101

102 :returns: Setting value as string, or ``None``.

103 """

104 sql = sa.text("select value from setting where name = :name")

105 return session.execute(sql, params={"name": name}).scalar()

106

107

108def make_engine_from_config(config_dict, prefix="sqlalchemy.", **kwargs):

109 """

110 Construct a new DB engine from configuration dict.

111

112 This is a wrapper around upstream

113 :func:`sqlalchemy:sqlalchemy.engine_from_config()`. For even

114 broader context of the SQLAlchemy

115 :class:`~sqlalchemy:sqlalchemy.engine.Engine` and their

116 configuration, see :doc:`sqlalchemy:core/engines`.

117

118 The purpose of the customization is to allow certain attributes of

119 the engine to be driven by config, whereas the upstream function

120 is more limited in that regard. The following in particular:

121

122 * ``poolclass``

123 * ``pool_pre_ping``

124

125 If these options are present in the configuration dict, they will

126 be coerced to appropriate Python equivalents and then passed as

127 kwargs to the upstream function.

128

129 An example config file leveraging this feature:

130

131 .. code-block:: ini

132

133 [wutta.db]

134 default.url = sqlite:///tmp/default.sqlite

135 default.poolclass = sqlalchemy.pool:NullPool

136 default.pool_pre_ping = true

137

138 Note that if present, the ``poolclass`` value must be a "spec"

139 string, as required by :func:`~wuttjamaican.util.load_object()`.

140 """

141 config_dict = dict(config_dict)

142

143 # convert 'poolclass' arg to actual class

144 key = f"{prefix}poolclass"

145 if key in config_dict and "poolclass" not in kwargs:

146 kwargs["poolclass"] = load_object(config_dict.pop(key))

147

148 # convert 'pool_pre_ping' arg to boolean

149 key = f"{prefix}pool_pre_ping"

150 if key in config_dict and "pool_pre_ping" not in kwargs:

151 kwargs["pool_pre_ping"] = parse_bool(config_dict.pop(key))

152

153 engine = sa.engine_from_config(config_dict, prefix, **kwargs)

154

155 return engine

156

157

158##############################

159# alembic functions

160##############################

161

162

163def make_alembic_config(config):

164 """

165 Make and return a new Alembic config object, based on current app

166 config.

167

168 This tries to set the following on the Alembic config:

169

170 * :attr:`~alembic:alembic.config.Config.config_file_name` - set to

171 app's primary config file

172 * main option ``script_location``

173 * main option ``version_locations``

174

175 The latter 2 are read normally from app config, then set on the

176 Alembic config via

177 :meth:`~alembic:alembic.config.Config.set_main_option()`.

178

179 .. note::

180

181 IIUC, Alembic should not need to attempt to read config values

182 from file, as long as we're able to set the above explicitly.

183 However we set the ``config_file_name`` "just in case" Alembic

184 needs it, but also to ensure it is discoverable from within the

185 ``env.py`` script...

186

187 When a migration script runs, code within ``env.py`` will call

188 :func:`make_config()` using the filename which it inspects from

189 the Alembic config.

190

191 (Confused yet?!)

192

193 :returns: :class:`alembic:alembic.config.Config` instance

194 """

195 alembic_config = AlembicConfig()

196

197 # TODO: not sure what we can do here besides assume the "primary"

198 # config file should be used?

199 if config.files_read:

200 alembic_config.config_file_name = config.get_prioritized_files()[0]

201

202 if script_location := config.get("alembic.script_location", usedb=False):

203 alembic_config.set_main_option("script_location", script_location)

204

205 if version_locations := config.get("alembic.version_locations", usedb=False):

206 alembic_config.set_main_option("version_locations", version_locations)

207

208 return alembic_config

209

210

211def get_alembic_scriptdir(config, alembic_config=None):

212 """

213 Get a "Script Directory" object for Alembic.

214

215 This allows for inspection of the migration scripts.

216

217 :param config: App config object.

218

219 :param alembic_config: Alembic config object, if you have one.

220 Otherwise :func:`make_alembic_config()` will be called.

221

222 :returns: :class:`~alembic:alembic.script.ScriptDirectory` instance

223 """

224 if not alembic_config:

225 alembic_config = make_alembic_config(config)

226 return ScriptDirectory.from_config(alembic_config)

227

228

229def check_alembic_current(config, alembic_config=None):

230 """

231 Compare the current revisions in the :term:`app database` to those

232 found in the migration scripts.

233

234 :param config: App config object.

235

236 :param alembic_config: Alembic config object, if you have one.

237 Otherwise :func:`make_alembic_config()` will be called.

238

239 :returns: ``True`` if the DB already has all migrations applied;

240 ``False`` if not.

241 """

242 script = get_alembic_scriptdir(config, alembic_config)

243 with config.appdb_engine.begin() as conn:

244 context = MigrationContext.configure(conn)

245 return set(context.get_current_heads()) == set(script.get_heads())