Coverage for python/felis/db/utils.py: 16%
178 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-24 01:22 -0800
1"""Database utility functions and classes."""
3# This file is part of felis.
4#
5# Developed for the LSST Data Management System.
6# This product includes software developed by the LSST Project
7# (https://www.lsst.org).
8# See the COPYRIGHT file at the top-level directory of this distribution
9# for details of code ownership.
10#
11# This program is free software: you can redistribute it and/or modify
12# it under the terms of the GNU General Public License as published by
13# the Free Software Foundation, either version 3 of the License, or
14# (at your option) any later version.
15#
16# This program is distributed in the hope that it will be useful,
17# but WITHOUT ANY WARRANTY; without even the implied warranty of
18# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
19# GNU General Public License for more details.
20#
21# You should have received a copy of the GNU General Public License
22# along with this program. If not, see <https://www.gnu.org/licenses/>.
24from __future__ import annotations
26import logging
27import re
28from typing import IO, Any
30from sqlalchemy import MetaData, inspect, types
31from sqlalchemy.engine import Connection, Dialect, Engine, ResultProxy
32from sqlalchemy.engine.mock import MockConnection, create_mock_engine
33from sqlalchemy.engine.url import URL
34from sqlalchemy.exc import SQLAlchemyError
35from sqlalchemy.schema import CreateSchema, DropSchema, Index, Table
36from sqlalchemy.sql import text
37from sqlalchemy.types import TypeEngine
39from .dialects import get_dialect_module
41__all__ = ["ConnectionWrapper", "DatabaseContext", "SQLWriter", "string_to_typeengine"]
43logger = logging.getLogger("felis")
45_DATATYPE_REGEXP = re.compile(r"(\w+)(\((.*)\))?")
46"""Regular expression to match data types with parameters in parentheses."""
49def string_to_typeengine(
50 type_string: str, dialect: Dialect | None = None, length: int | None = None
51) -> TypeEngine:
52 """Convert a string representation of a datatype to a SQLAlchemy type.
54 Parameters
55 ----------
56 type_string
57 The string representation of the data type.
58 dialect
59 The SQLAlchemy dialect to use. If None, the default dialect will be
60 used.
61 length
62 The length of the data type. If the data type does not have a length
63 attribute, this parameter will be ignored.
65 Returns
66 -------
67 `sqlalchemy.types.TypeEngine`
68 The SQLAlchemy type engine object.
70 Raises
71 ------
72 ValueError
73 Raised if the type string is invalid or the type is not supported.
75 Notes
76 -----
77 This function is used when converting type override strings defined in
78 fields such as ``mysql:datatype`` in the schema data.
79 """
80 match = _DATATYPE_REGEXP.search(type_string)
81 if not match:
82 raise ValueError(f"Invalid type string: {type_string}")
84 type_name, _, params = match.groups()
85 if dialect is None:
86 type_class = getattr(types, type_name.upper(), None)
87 else:
88 try:
89 dialect_module = get_dialect_module(dialect.name)
90 except KeyError:
91 raise ValueError(f"Unsupported dialect: {dialect}")
92 type_class = getattr(dialect_module, type_name.upper(), None)
94 if not type_class:
95 raise ValueError(f"Unsupported type: {type_class}")
97 if params:
98 params = [int(param) if param.isdigit() else param for param in params.split(",")]
99 type_obj = type_class(*params)
100 else:
101 type_obj = type_class()
103 if hasattr(type_obj, "length") and getattr(type_obj, "length") is None and length is not None:
104 type_obj.length = length
106 return type_obj
109def is_mock_url(url: URL) -> bool:
110 """Check if the engine URL is a mock URL.
112 Parameters
113 ----------
114 url
115 The SQLAlchemy engine URL.
117 Returns
118 -------
119 bool
120 True if the URL is a mock URL, False otherwise.
121 """
122 return (url.drivername == "sqlite" and url.database is None) or (
123 url.drivername != "sqlite" and url.host is None
124 )
127def is_valid_engine(engine: Engine | MockConnection | None) -> bool:
128 """Check if the engine is valid.
130 The engine cannot be none; it must not be a mock connection; and it must
131 not be a mock URL which is missing a host or, for sqlite, a database name.
133 Parameters
134 ----------
135 engine
136 The SQLAlchemy engine or mock connection.
138 Returns
139 -------
140 bool
141 True if the engine is valid, False otherwise.
142 """
143 return engine is not None and not isinstance(engine, MockConnection) and not is_mock_url(engine.url)
146class SQLWriter:
147 """Write SQL statements to stdout or a file.
149 Parameters
150 ----------
151 file
152 The file to write the SQL statements to. If None, the statements
153 will be written to stdout.
154 """
156 def __init__(self, file: IO[str] | None = None) -> None:
157 """Initialize the SQL writer."""
158 self.file = file
159 self.dialect: Dialect | None = None
161 def write(self, sql: Any, *multiparams: Any, **params: Any) -> None:
162 """Write the SQL statement to a file or stdout.
164 Statements with parameters will be formatted with the values
165 inserted into the resultant SQL output.
167 Parameters
168 ----------
169 sql
170 The SQL statement to write.
171 *multiparams
172 The multiparams to use for the SQL statement.
173 **params
174 The params to use for the SQL statement.
176 Notes
177 -----
178 The functions arguments are typed very loosely because this method in
179 SQLAlchemy is untyped, amd we do not call it directly.
180 """
181 compiled = sql.compile(dialect=self.dialect)
182 sql_str = str(compiled) + ";"
183 params_list = [compiled.params]
184 for params in params_list:
185 if not params:
186 print(sql_str, file=self.file)
187 continue
188 new_params = {}
189 for key, value in params.items():
190 if isinstance(value, str):
191 new_params[key] = f"'{value}'"
192 elif value is None:
193 new_params[key] = "null"
194 else:
195 new_params[key] = value
196 print(sql_str % new_params, file=self.file)
199class ConnectionWrapper:
200 """Wrap a SQLAlchemy engine or mock connection to provide a consistent
201 interface for executing SQL statements.
203 Parameters
204 ----------
205 engine
206 The SQLAlchemy engine or mock connection to wrap.
207 """
209 def __init__(self, engine: Engine | MockConnection):
210 """Initialize the connection wrapper."""
211 self.engine = engine
213 def execute(self, statement: Any) -> ResultProxy:
214 """Execute a SQL statement on the engine and return the result.
216 Parameters
217 ----------
218 statement
219 The SQL statement to execute.
221 Returns
222 -------
223 ``sqlalchemy.engine.ResultProxy``
224 The result of the statement execution.
226 Notes
227 -----
228 The statement will be executed in a transaction block if not using
229 a mock connection.
230 """
231 if isinstance(statement, str):
232 statement = text(statement)
233 if isinstance(self.engine, Engine):
234 try:
235 with self.engine.begin() as connection:
236 result = connection.execute(statement)
237 return result
238 except SQLAlchemyError as e:
239 connection.rollback()
240 logger.error(f"Error executing statement: {e}")
241 raise
242 elif isinstance(self.engine, MockConnection):
243 return self.engine.connect().execute(statement)
244 else:
245 raise ValueError("Unsupported engine type:" + str(type(self.engine)))
248class DatabaseContext:
249 """Manage the database connection and SQLAlchemy metadata.
251 Parameters
252 ----------
253 metadata
254 The SQLAlchemy metadata object.
256 engine
257 The SQLAlchemy engine or mock connection object.
258 """
260 def __init__(self, metadata: MetaData, engine: Engine | MockConnection):
261 """Initialize the database context."""
262 self.engine = engine
263 self.dialect_name = engine.dialect.name
264 self.metadata = metadata
265 self.connection = ConnectionWrapper(engine)
267 def initialize(self) -> None:
268 """Create the schema in the database if it does not exist.
270 Raises
271 ------
272 ValueError
273 Raised if the database is not supported or it already exists.
274 sqlalchemy.exc.SQLAlchemyError
275 Raised if there is an error creating the schema.
277 Notes
278 -----
279 In MySQL, this will create a new database and, in PostgreSQL, it will
280 create a new schema. For other variants, this is an unsupported
281 operation.
282 """
283 if self.engine.dialect.name == "sqlite":
284 # Initialization is unneeded for sqlite.
285 return
286 schema_name = self.metadata.schema
287 if schema_name is None:
288 raise ValueError("Schema name is required to initialize the schema.")
289 try:
290 if self.dialect_name == "mysql":
291 logger.debug(f"Checking if MySQL database exists: {schema_name}")
292 result = self.execute(text(f"SHOW DATABASES LIKE '{schema_name}'"))
293 if result.fetchone():
294 raise ValueError(f"MySQL database '{schema_name}' already exists.")
295 logger.debug(f"Creating MySQL database: {schema_name}")
296 self.execute(text(f"CREATE DATABASE {schema_name}"))
297 elif self.dialect_name == "postgresql":
298 logger.debug(f"Checking if PG schema exists: {schema_name}")
299 result = self.execute(
300 text(
301 f"""
302 SELECT schema_name
303 FROM information_schema.schemata
304 WHERE schema_name = '{schema_name}'
305 """
306 )
307 )
308 if result.fetchone():
309 raise ValueError(f"PostgreSQL schema '{schema_name}' already exists.")
310 logger.debug(f"Creating PG schema: {schema_name}")
311 self.execute(CreateSchema(schema_name))
312 else:
313 raise ValueError(f"Initialization not supported for: {self.dialect_name}")
314 except SQLAlchemyError as e:
315 logger.error(f"Error creating schema: {e}")
316 raise
318 def drop(self) -> None:
319 """Drop the schema in the database if it exists.
321 Raises
322 ------
323 ValueError
324 Raised if the database is not supported.
326 Notes
327 -----
328 In MySQL, this will drop a database. In PostgreSQL, it will drop a
329 schema. A SQlite database will have all its tables dropped. For other
330 database variants, this is currently an unsupported operation.
331 """
332 try:
333 if self.dialect_name == "sqlite":
334 if isinstance(self.engine, Engine):
335 logger.debug("Dropping tables in SQLite schema")
336 self.metadata.drop_all(bind=self.engine)
337 else:
338 schema_name = self.metadata.schema
339 if schema_name is None:
340 raise ValueError("Schema name is required to drop the schema.")
341 if self.dialect_name == "mysql":
342 logger.debug(f"Dropping MySQL database if exists: {schema_name}")
343 self.execute(text(f"DROP DATABASE IF EXISTS {schema_name}"))
344 elif self.dialect_name == "postgresql":
345 logger.debug(f"Dropping PostgreSQL schema if exists: {schema_name}")
346 self.execute(DropSchema(schema_name, if_exists=True, cascade=True))
347 except SQLAlchemyError as e:
348 logger.error(f"Error dropping schema: {e}")
349 raise
351 def create_all(self) -> None:
352 """Create all tables in the schema using the metadata object."""
353 if isinstance(self.engine, Engine):
354 # Use a transaction for a real connection.
355 with self.engine.begin() as conn:
356 try:
357 self.metadata.create_all(bind=conn)
358 conn.commit()
359 except SQLAlchemyError as e:
360 conn.rollback()
361 logger.error(f"Error creating tables: {e}")
362 raise
363 elif isinstance(self.engine, MockConnection):
364 # Mock connection so no need for a transaction.
365 self.metadata.create_all(self.engine)
366 else:
367 raise ValueError("Unsupported engine type: " + str(type(self.engine)))
369 def create_indexes(self) -> None:
370 """Create all indexes in the schema using the metadata object."""
371 assert isinstance(self.engine, Engine)
372 with self.engine.begin() as conn:
373 for table in self.metadata.tables.values():
374 for index in table.indexes:
375 if self._index_exists(conn, table, index):
376 logger.warning(f"Skipping creation of index '{index.name}' which already exists")
377 continue
378 index.create(bind=conn, checkfirst=True)
379 logger.info(f"Created index '{index.name}'")
381 def drop_indexes(self) -> None:
382 """Drop all indexes in the schema using the metadata object."""
383 assert isinstance(self.engine, Engine)
384 with self.engine.begin() as conn:
385 for table in self.metadata.tables.values():
386 for index in table.indexes:
387 if not self._index_exists(conn, table, index):
388 logger.warning(f"Skipping index '{index.name}' which does not exist")
389 continue
390 index.drop(bind=conn, checkfirst=True)
391 logger.info(f"Dropped index '{index.name}'")
393 @classmethod
394 def _index_exists(cls, conn: Connection, table: Table, index: Index) -> bool:
395 """Return True if an index with this name exists on the table in the
396 database.
397 """
398 if index.name is None:
399 # Anonymous indexes can't be reliably checked by name.
400 return False
402 inspector = inspect(conn)
403 # For SQLite, schema should be None as SQLite doesn't support schemas
404 # even though the table.schema might be set to "main"
405 schema = None if conn.dialect.name == "sqlite" else table.schema
406 existing = {
407 ix["name"]
408 for ix in inspector.get_indexes(
409 table_name=table.name,
410 schema=schema,
411 )
412 if "name" in ix and ix["name"] is not None
413 }
414 return index.name in existing
416 @staticmethod
417 def create_mock_engine(engine_url: str | URL, output_file: IO[str] | None = None) -> MockConnection:
418 """Create a mock engine for testing or dumping DDL statements.
420 Parameters
421 ----------
422 engine_url
423 The SQLAlchemy engine URL.
424 output_file
425 The file to write the SQL statements to. If None, the statements
426 will be written to stdout.
428 Returns
429 -------
430 ``sqlalchemy.engine.mock.MockConnection``
431 The mock connection object.
432 """
433 writer = SQLWriter(output_file)
434 engine = create_mock_engine(engine_url, executor=writer.write, paramstyle="pyformat")
435 writer.dialect = engine.dialect
436 return engine
438 def execute(self, statement: Any) -> ResultProxy:
439 """Execute a SQL statement on the engine and return the result.
441 Parameters
442 ----------
443 statement
444 The SQL statement to execute.
446 Returns
447 -------
448 ``sqlalchemy.engine.ResultProxy``
449 The result of the statement execution.
451 Notes
452 -----
453 This is just a wrapper around the execution method of the connection
454 object, which may execute on a real or mock connection.
455 """
456 return self.connection.execute(statement)