Coverage for src/mafw/processor_library/db_init.py: 100%
172 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-26 09:13 +0000
1# Copyright 2025–2026 European Union
2# Author: Bulgheroni Antonio (antonio.bulgheroni@ec.europa.eu)
3# SPDX-License-Identifier: EUPL-1.2
4"""
5Database initialisation processor module.
7This module contains the following processors:
9:class:`TableCreator`
10 processor which handles the creation of database tables based on registered models. It provides functionality to
11 create tables automatically while respecting existing tables and offering options for forced recreation.
13:class:`TriggerRefresher`
14 processor to safely update the trigger definitions. It removes all existing triggers and regenerates them
15 according to the new definition. Particularly useful when debugging triggers, it can also be left at the
16 beginning of all analysis pipelines since it does not cause any loss of data.
18:class:`SQLScriptRunner`
19 processor to execute SQL scripts from files against the database. It reads SQL files, removes block comments,
20 splits the content into individual statements, and executes them within a transaction.
22.. versionadded:: v2.0.0
23"""
25import logging
26import re
27from collections.abc import Collection
28from pathlib import Path
29from typing import TYPE_CHECKING, Any, cast
31import peewee
32from rich.prompt import Confirm, Prompt
34from mafw.db.db_model import mafw_model_register
35from mafw.db.db_types import PeeweeModelWithMeta
36from mafw.db.std_tables import StandardTable
37from mafw.db.trigger import MySQLDialect, PostgreSQLDialect, SQLiteDialect, TriggerDialect
38from mafw.decorators import database_required, single_loop
39from mafw.enumerators import LoopingStatus, ProcessorExitStatus
40from mafw.mafw_errors import InvalidConfigurationError, UnsupportedDatabaseError
41from mafw.processor import ActiveParameter, Processor
43log = logging.getLogger(__name__)
45block_comment_re = re.compile(r'/\*.*?\*/', re.DOTALL)
48def _standard_table_models() -> set[type[StandardTable]]:
49 """Return the registered standard-table models when available."""
50 standard_tables = mafw_model_register.get_standard_tables()
51 if isinstance(standard_tables, list):
52 return set(standard_tables)
53 return set()
56@database_required
57@single_loop
58class TableCreator(Processor):
59 """
60 Processor to create all tables in the database.
62 This processor can be included in all pipelines in order to create all tables in the database. Its functionality
63 is based on the fact that all :class:`.MAFwBaseModel` subclasses are automatically included in a global register
64 (:data:`.mafw_model_register`).
66 This processor will perform the following:
68 #. Get a list of all tables already existing in the database.
69 #. Prune from the lists of models the ones for which already exist in the database.
70 #. Create the remaining tables.
73 This overall behaviour can be modified via the following parameters:
75 * *force_recreate* (bool, default = False): Use with extreme care. When set to True, all tables in the
76 database and in the model register will be first dropped and then recreated. It is almost equivalent to a re-initialization of the
77 whole DB with all the data being lost.
79 * *soft_recreate* (bool, default = True): When set to true, all tables whose model is in the mafw model
80 register will be recreated with the safe flag. It means that there won't be any table drop. If a table is
81 already existing, nothing will happen. If a new trigger is added to the table this will be created. When
82 set to False, only tables whose model is in the register and that are not existing will be created.
84 * *apply_only_to_prefix* (list[str], default = []): This parameter allows to create only the tables that do
85 not already exist and whose name start with one of the provided prefixes.
87 .. versionadded:: v2.0.0
89 """
91 force_recreate = ActiveParameter(
92 name='force_recreate', default=False, help_doc='First drop and then create the tables. LOSS OF ALL DATA!!!'
93 )
94 """
95 Force recreate (bool, default = False).
97 Use with extreme care. When set to True, all tables in the database and in the model register will be first
98 dropped and then recreated. It is almost equivalent to a re-initialization of the whole DB with all the data
99 being lost.
100 """
102 soft_recreate = ActiveParameter(
103 name='soft_recreate', default=True, help_doc='Safe recreate tables without dropping. No data loss'
104 )
105 """
106 Soft recreate (bool default = True).
108 When set to true, all tables whose model is in the mafw model register will be recreated with the safe flag. It
109 means that there won't be any table drop. If a table is already existing, nothing will happen. If a new trigger
110 is added to the table, this will be created. When set to False, only tables whose model is in the register and
111 that are not existing will be created.
112 """
114 apply_only_to_prefix = ActiveParameter[list[str]](
115 name='apply_only_to_prefix',
116 default=[],
117 help_doc='Create only tables whose name start with the provided prefixes.',
118 )
119 """
120 Apply only to tables starting with prefix (list[str], default = []).
122 This parameter allows to create only the tables that do not already exist and whose name start with one of the
123 provided prefixes.
124 """
126 def __init__(self, *args: Any, **kwargs: Any) -> None:
127 super().__init__(*args, **kwargs)
128 self.existing_table_names: list[str] = []
129 """The list of all existing tables in the database."""
131 def validate_configuration(self) -> None:
132 """
133 Configuration validation
135 :attr:`force_recreate` and :attr:`soft_recreate` cannot be both valid.
137 :raises InvalidConfigurationError: if both recreate types are True.
138 """
139 if self.force_recreate and self.soft_recreate:
140 raise InvalidConfigurationError(
141 'Both force_recreate and soft_recreate set to True. Incompatible configuration'
142 )
144 def process(self) -> None:
145 """
146 Execute the table creation process.
148 This method performs the following steps:
150 #. Identify all models that have automatic creation enabled.
151 #. Filter models based on the apply_only_to_prefix parameter if specified.
152 #. Handle forced recreation if requested, including user confirmation.
153 #. Handle soft recreation if requested, letting all tables with a known model be recreated.
154 #. Create the required tables.
155 #. Initialise standard tables after recreation if needed.
157 If user cancel the creation, the processor exit status is set to :attr:`.ProcessorExitStatus.Aborted` so that
158 the whole processor list is blocked.
159 """
160 # get all tables with autocreation flag
161 autocreation_models = [
162 model
163 for _, model in mafw_model_register.items()
164 if model._meta.automatic_creation # type: ignore[attr-defined]
165 ]
167 # filter out standard tables if the user decided not to create them
168 standard_table_models = _standard_table_models()
169 if not self.create_standard_tables:
170 autocreation_models = [model for model in autocreation_models if model not in standard_table_models]
172 # get the table name from the model class
173 autocreation_table_names = [cast(PeeweeModelWithMeta, model)._meta.table_name for model in autocreation_models]
175 if self.apply_only_to_prefix:
176 # remove tables with all given prefixes
177 autocreation_table_names = [
178 name
179 for name in autocreation_table_names
180 if any([name.startswith(prefix) for prefix in self.apply_only_to_prefix]) # type: ignore[union-attr]
181 ]
183 # in the case of force_recreation, we need to have user confirmation
184 if self.force_recreate:
185 log.warning(f'Forcing recreation of {len(autocreation_table_names)} tables in the database.')
186 log.warning('All data in these tables will be lost.')
188 with self._user_interface.enter_interactive_mode():
189 question = 'Are you really sure?'
190 if self._user_interface.name == 'rich':
191 question = '[red][bold]' + question + '[/red][/bold]'
192 confirmation = self._user_interface.prompt_question(
193 question=question, prompt_type=Confirm, default=False, show_default=True, case_sensitive=True
194 )
195 if not confirmation:
196 self.processor_exit_status = ProcessorExitStatus.Aborted
197 return
198 else:
199 log.info(f'Removing {len(autocreation_table_names)} tables from the database.')
200 models = [
201 model
202 for model in autocreation_models
203 if cast(PeeweeModelWithMeta, model)._meta.table_name in autocreation_table_names
204 ]
205 self.database.drop_tables(models) # type: ignore[arg-type]
207 if self.soft_recreate:
208 # recreate all tables in the mafw register
209 models = [
210 model
211 for model in autocreation_models
212 if cast(PeeweeModelWithMeta, model)._meta.table_name in autocreation_table_names
213 ]
214 else:
215 # recreate all tables in the mafw register and that are not yet existing
216 self.existing_table_names = self.database.get_tables()
217 models = [
218 model
219 for model in autocreation_models
220 if cast(PeeweeModelWithMeta, model)._meta.table_name in autocreation_table_names
221 and cast(PeeweeModelWithMeta, model)._meta.table_name not in self.existing_table_names
222 ]
223 self.database.create_tables(models) # type: ignore[arg-type]
225 if self.force_recreate and self.create_standard_tables:
226 # in the case of a recreation, do the initialisation of all dropped standard tables.
227 for model in models:
228 if model in standard_table_models:
229 model.init()
231 n = len(models)
232 if n > 0:
233 if n == 1:
234 plu = ''
235 else:
236 plu = 's'
237 if self.soft_recreate:
238 soft = '(soft) '
239 else:
240 soft = ''
242 log.info(f'Successfully {soft}created {len(models)} table{plu}.')
245@database_required
246class TriggerRefresher(Processor):
247 """
248 Processor to recreate all triggers.
250 Triggers are database objects, and even though they could be created, dropped and modified at any moment,
251 within the MAFw execution cycle they are normally created along with the table they are targeting.
253 When the table is created, also all its triggers are created,
254 but unless differently specified, with the safe flag on, that means that they are created if they do not exist.
256 This might be particularly annoying when modifying an existing trigger, because you need to manually drop the
257 trigger to let the table creation mechanism to create the newer version.
259 The goal of this processor is to drop all existing triggers and then recreate the corresponding tables so to have
260 an updated version of the triggers.
262 The processor is relying on the fact that all subclasses of :class:`.MAFwBaseModel`
263 are automatically inserted in the :data:`.mafw_model_register` so that the model class can be retrieved from the
264 table name.
266 Before removing any trigger, the processor will build a list with all the affected tables and check if all of
267 them are in the :data:`.mafw_model_register`, if so, it will proceed without asking any further confirmation.
268 Otherwise, if some affected tables are not in the register, then it will ask the user to decide what to do:
270 - Remove only the triggers whose tables are in the register and thus recreated afterward.
271 - Remove all triggers, in this case, some of them will not be recreated.
272 - Abort the processor.
274 Trigger manipulations (drop and creation) are not directly implemented in :link:`peewee` and are an extension
275 provided by MAFw. In order to be compatible with the three main databases (:link:`sqlite`, :link:`mysql` and
276 :link:`postgresql`), the SQL generation is obtained via the :class:`.TriggerDialect` interface.
278 .. seealso::
280 The :class:`.Trigger` class and also the :ref:`trigger chapter <triggers>` for a deeper explanation on triggers.
282 The :class:`.ModelRegister` class, the :data:`.mafw_model_register` and the :ref:`related chapter
283 <auto_registration>` on the automatic registration mechanism.
285 The :class:`.TriggerDialect` and its subclasses, for a database independent way to generate SQL statement
286 related to triggers.
288 .. versionadded:: v2.0.0
289 """
291 def __init__(self, *args: Any, **kwargs: Any) -> None:
292 super().__init__(*args, **kwargs)
293 self.dialect: TriggerDialect | None = None
294 self.tables_to_be_rebuilt: set[str] = set()
296 def get_dialect(self) -> TriggerDialect:
297 """
298 Get the valid SQL dialect based on the type of Database
300 :return: The SQL trigger dialect
301 :type: :class:`.TriggerDialect`
302 :raises: :class:`.UnsupportedDatabaseError` if there is no dialect for the current DB.
303 """
304 if self.dialect is not None:
305 return self.dialect
307 if self._database is None:
308 # Default to SQLite dialect
309 return SQLiteDialect()
311 db = self._database
312 if isinstance(db, peewee.DatabaseProxy):
313 db = db.obj # Get the actual database from the proxy
315 dialect: TriggerDialect
316 if isinstance(db, peewee.SqliteDatabase):
317 dialect = SQLiteDialect()
318 elif isinstance(db, peewee.MySQLDatabase):
319 dialect = MySQLDialect()
320 elif isinstance(db, peewee.PostgresqlDatabase):
321 dialect = PostgreSQLDialect()
322 else:
323 raise UnsupportedDatabaseError(f'Unsupported database type: {type(db)}')
325 return dialect
327 def start(self) -> None:
328 super().start()
329 self.dialect = self.get_dialect()
331 def get_items(self) -> Collection[Any]:
332 """
333 Retrieves a list of database triggers and interacts with the user to determine which ones to process.
335 This method fetches all currently defined database triggers. If any tables
336 associated with these triggers are not known (i.e., not registered in
337 :data:`.mafw_model_register`), it enters an interactive mode to prompt the user for
338 a course of action:
340 1. **Remove All Triggers (A):** Processes all triggers for subsequent removal,
341 but only marks 'rebuildable' tables for rebuilding.
342 2. **Remove Only Rebuildable Triggers (O):** Processes only triggers associated
343 with 'rebuildable' tables.
344 3. **Quit (Q):** Aborts the entire process.
346 If no unknown tables are found, or the user chooses to process rebuildable tables,
347 the list of triggers and the set of tables to be rebuilt are prepared for the next stage.
349 :return: A collection of database triggers to be processed, in the for tuple trigger_name, table_name
350 :rtype: list[tuple[str, str]]
351 """
352 if TYPE_CHECKING:
353 assert self.dialect is not None
355 s: list[tuple[str, str]] = self.database.execute_sql(self.dialect.select_all_trigger_sql()).fetchall() # type: ignore[no-untyped-call]
356 tables = [r[1] for r in s]
358 affected_tables = set(tables)
359 known_tables = mafw_model_register.get_table_names()
360 rebuildable_tables = {t for t in affected_tables if t in known_tables}
361 not_rebuildable_tables = affected_tables - rebuildable_tables
363 if len(not_rebuildable_tables) > 0:
364 log.warning(f'There are some tables ({len(not_rebuildable_tables)}) that cannot be rebuild')
365 with self._user_interface.enter_interactive_mode():
366 question = 'Remove all triggers (A), remove only rebuildable triggers (O), quit (Q)'
367 if self._user_interface.name == 'rich':
368 question = '[red][bold]' + question + '[/red][/bold]'
370 class TriggerPrompt(Prompt):
371 response_type = str
372 validate_error_message = '[prompt.invalid]Please enter A, O or Q'
373 choices: list[str] = ['A', 'O', 'Q']
375 answer = self._user_interface.prompt_question(
376 question=question,
377 prompt_type=TriggerPrompt,
378 default='O',
379 show_default=True,
380 case_sensitive=False,
381 show_answer=True,
382 )
384 if answer == 'Q':
385 s = []
386 affected_tables = set()
387 self.processor_exit_status = ProcessorExitStatus.Aborted
388 self.looping_status = LoopingStatus.Abort
390 elif answer == 'O':
391 s = [r for r in s if r[1] in rebuildable_tables]
392 affected_tables = rebuildable_tables
394 else: # equivalent to 'A'
395 # remove all triggers
396 # but rebuilds only rebuildable_tables
397 affected_tables = rebuildable_tables
399 self.tables_to_be_rebuilt = affected_tables
400 return s
402 def process(self) -> None:
403 """Delete the current trigger from its table"""
404 if TYPE_CHECKING:
405 assert self.dialect is not None
407 self.database.execute_sql(self.dialect.drop_trigger_sql(self.item[0], safe=True, table_name=self.item[1])) # type: ignore[no-untyped-call]
409 def finish(self) -> None:
410 """
411 Recreate the tables from which triggers were dropped.
413 This is only done if the user did not abort the process.
414 """
415 if self.looping_status != LoopingStatus.Abort:
416 log.info(f'Recreating {self.n_item} triggers on {len(self.tables_to_be_rebuilt)} tables...')
417 models = [mafw_model_register.get_model(table_name) for table_name in self.tables_to_be_rebuilt]
419 self.database.create_tables(models) # type: ignore[arg-type]
420 super().finish()
422 def format_progress_message(self) -> None:
423 self.progress_message = f'Dropping trigger {self.item[0]} from table {self.item[1]}'
426@database_required
427class SQLScriptRunner(Processor):
428 """
429 Processor to execute SQL scripts from files against the database.
431 This processor reads SQL files, removes multi-line block comments, splits the content into individual
432 statements, and executes them within a transaction. It is designed to handle SQL script execution
433 in a safe manner by wrapping all statements in a single atomic transaction.
435 The processor accepts a list of SQL files through the :attr:`sql_files` parameter. Each file is validated
436 to ensure it exists and is a regular file before processing. Block comments (`/* ... */`) are removed
437 from the SQL content before statement parsing.
439 .. versionadded:: v2.0.0
440 """
442 sql_files = ActiveParameter[list[Path]]('sql_files', default=[], help_doc='A list of SQL files to be processed')
443 """List of SQL files to be processed"""
445 def validate_configuration(self) -> None:
446 """
447 Validate the configuration of SQL script runner.
449 Ensures that all specified SQL files exist and are regular files.
451 :raises InvalidConfigurationError: if any of the specified files does not exist or is not a regular file.
452 """
453 if TYPE_CHECKING:
454 # we need to convince mypy that the sql_files is not and ActiveParameter but
455 # the content of the ActiveParameter
456 assert isinstance(self.sql_files, list)
458 self.sql_files = [Path(file) for file in self.sql_files]
459 for file in self.sql_files:
460 if not file.exists() or not file.is_file():
461 raise InvalidConfigurationError(f'There are issues with SQL file "{file.resolve()}". Please verify.')
463 def get_items(self) -> Collection[Any]:
464 """
465 Get the collection of SQL files to be processed.
467 :return: A collection of SQL file paths to be processed
468 :rtype: Collection[Any]
469 """
470 if TYPE_CHECKING:
471 # we need to convince mypy that the sql_files is not and ActiveParameter but
472 # the content of the ActiveParameter
473 assert isinstance(self.sql_files, list)
475 return self.sql_files
477 def process(self) -> None:
478 """
479 Process a single SQL file by reading, parsing, and executing its statements.
481 Reads the SQL file content, removes multi-line block comments, splits the content
482 into individual SQL statements, and executes them within a transaction.
484 If no statements are found in the file, a warning is logged. If an error occurs
485 during execution, the transaction is rolled back and the exception is re-raised.
487 :raises Exception: If an error occurs during SQL statement execution.
488 """
489 with open(self.item) as sql_file:
490 sql_content = sql_file.read()
492 # remove the multi-line block comments (/* ... */)
493 sql_content = block_comment_re.sub('', sql_content)
495 statements = [s.strip() + ';' for s in sql_content.split(';') if s.strip()]
497 if not statements:
498 log.warning(f'No SQL statements found to execute in {self.item.name}.')
500 log.debug(f'Found {len(statements)} statements to execute.')
502 try:
503 # use an atomic transaction to wrap the execution of all statements.
504 with self.database.atomic():
505 for statement in statements:
506 self.database.execute_sql(statement) # type: ignore[no-untyped-call]
508 except Exception as e:
509 log.critical(f'An error occurred while executing the SQL script {self.item.name}.')
510 log.critical('Rolling back the database to preserve integrity.')
511 log.critical(f'Error details: {e}')
512 raise
514 def format_progress_message(self) -> None:
515 self.progress_message = f'Processing SQL file {self.item.name}'