Coverage for src / mafw / db / model_register.py: 100%
74 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-09 09:08 +0000
1# Copyright 2025 European Union
2# Author: Bulgheroni Antonio (antonio.bulgheroni@ec.europa.eu)
3# SPDX-License-Identifier: EUPL-1.2
4"""
5Module for registering and retrieving database models.
7This module provides functionality to register database models with table names as key
8and retrieve them using flexible naming conventions that support prefixes and suffixes.
10.. versionadded:: v2.0.0
11"""
13import itertools
14import logging
15import threading
16from typing import TYPE_CHECKING, List, Type
18import peewee
20if TYPE_CHECKING:
21 from mafw.db.std_tables import StandardTable
23log = logging.getLogger(__name__)
26class ModelRegister:
27 """
28 A registry for database models with support for prefixes and suffixes.
30 This class allows registration of database models with table names and
31 provides flexible retrieval mechanisms that can handle different naming conventions.
33 The following implemented methods are thread-safe:
34 - :meth:`.register_model`
35 - :meth:`.register_prefix`
36 - :meth:`.register_suffix`
37 - :meth:`.get_model`
38 - :meth:`.get_model_names`
39 - :meth:`.get_table_names`
40 - :meth:`.items`
42 Accessing the class underling containers is instead **not** thread safe.
45 .. versionadded:: v2.0.0
46 """
48 def __init__(self) -> None:
49 self.models: dict[str, peewee.ModelBase] = {}
50 self.model_names: dict[str, str] = {} # this dictionary goes from ModelName to table_name
51 self.prefixes: list[str] = []
52 self.suffixes: list[str] = []
53 self._lock = threading.Lock()
55 def register_model(self, table_name: str, model: peewee.ModelBase) -> None:
56 """
57 Register a model with a specific table name.
59 If a model with the same table name already exists, it will be replaced
60 with a warning message.
62 :param table_name: The table name to register the model under
63 :type table_name: str
64 :param model: The peewee Model class to register
65 :type model: peewee.Model
66 """
67 with self._lock:
68 if table_name in self.models:
69 log.warning(f'A model with the same name ({table_name}) already exists. Replacing it.')
70 self.models[table_name] = model
71 self.model_names[model.__name__] = table_name
73 def register_prefix(self, prefix: str) -> None:
74 """
75 Register a prefix to be used when searching for models.
77 :param prefix: The prefix string to register
78 :type prefix: str
79 """
80 with self._lock:
81 if prefix not in self.prefixes:
82 self.prefixes.append(prefix)
84 def register_suffix(self, suffix: str) -> None:
85 """
86 Register a suffix to be used when searching for models.
88 :param suffix: The suffix string to register
89 :type suffix: str
90 """
91 with self._lock:
92 if suffix not in self.suffixes:
93 self.suffixes.append(suffix)
95 def get_model(self, name: str) -> peewee.ModelBase:
96 """
97 Retrieve a model by name, supporting prefixes and suffixes.
99 `name` could be either the table_name or the ModelName.
101 This method attempts to find a model by the given name, considering
102 registered prefixes and suffixes. It also handles conversion between
103 CamelCase model names and snake_case table names using peewee's utility.
105 :param name: The name to search for
106 :type name: str
107 :return: The registered peewee Model class
108 :rtype: peewee.Model
109 :raises KeyError: If no matching model is found or multiple similar models exist
110 """
111 with self._lock:
112 if name in self.models:
113 # let's assume we got a table_name
114 return self.models[name]
115 elif name in self.model_names:
116 # let's try with a ModelName
117 return self.models[self.model_names[name]]
118 else:
119 # let's try some combinations as a last resort!
120 prefixes = ['']
121 prefixes.extend(self.prefixes)
123 names = [name]
124 if peewee.make_snake_case(name) not in names: # type: ignore[attr-defined]
125 names.append(peewee.make_snake_case(name)) # type: ignore[attr-defined]
127 suffixes = ['']
128 suffixes.extend(self.suffixes)
130 combinations = itertools.product(prefixes, names, suffixes)
132 possible_names = []
133 for p, n, s in combinations:
134 name_under_test = p + n + s
135 if name_under_test in self.models:
136 possible_names.append(name_under_test)
137 possible_names = list(set(possible_names))
139 if len(possible_names) == 0:
140 log.error(f'Model {name} not found in the registered models')
141 log.error(f'Available models: {", ".join(self.models.keys())}')
142 raise KeyError(f'Model {name} not registered')
143 elif len(possible_names) == 1:
144 log.debug(f'Model {name} not found, but {possible_names[0]} is available. Using this model.')
145 return self.models[possible_names[0]]
146 else:
147 log.error(f'Model {name} not found in the registered models')
148 log.error(f'Following similar models found: {", ".join(possible_names)}')
149 raise KeyError(f'Model {name} not registered, but multiple similar ones found.')
151 def get_model_names(self) -> list[str]:
152 """
153 Get a list of all registered model names.
155 :return: List of registered model names
156 :rtype: list[str]
157 """
158 with self._lock:
159 return list(self.model_names.keys())
161 def get_table_names(self) -> list[str]:
162 """
163 Get a list of all registered table names.
165 :return: List of registered table names
166 :rtype: list[str]
167 """
168 with self._lock:
169 return list(self.models.keys())
171 def items(self) -> list[tuple[str, peewee.ModelBase]]:
172 """
173 Return the items list of the registered models dictionary.
175 This method provides access to all registered models through a dictionary-like
176 items view, allowing iteration over key-value pairs of table names and their
177 corresponding model classes.
179 In order to release the thread lock as soon as possible, instead of providing an iterator a list of the
180 current snapshot of the dictionary is provided.
182 :return: An items view of the registered models
183 :rtype: list[[str, peewee.ModelBase]]
184 """
185 with self._lock:
186 return list(self.models.items())
188 def get_standard_tables(self) -> List[Type['StandardTable']]:
189 """
190 Retrieve all registered models that are instances of :class:`~mafw.db.std_tables.StandardTable`.
192 This method filters the registered models and returns only those that inherit from
193 the :class:`~mafw.db.std_tables.StandardTable` base class.
195 This is useful for identifying and working with standard database tables that follow a specific
196 structure or interface.
198 Since the introduction of the :class:`~.ModelRegister`, there is no need any more for a standard table
199 plugin hook, instead the user can use this method to retrieve all standard tables.
201 :return: A list of registered model classes that are standard tables
202 :rtype: list[peewee.ModelBase]
203 """
204 from mafw.db.std_tables import StandardTable
206 return [t for t in self.models.values() if issubclass(t, StandardTable)]
208 def clear(self) -> None:
209 """
210 Clear all registered models, prefixes, and suffixes from the registry.
212 This method removes all entries from the internal dictionaries and lists,
213 effectively resetting the ModelRegister to its initial empty state.
215 .. note::
216 This operation cannot be undone. All previously registered models
217 and naming conventions will be lost after calling this method.
219 :rtype: None
220 """
221 self.models = {}
222 self.model_names = {}
223 self.prefixes = []
224 self.suffixes = []