]> git.openstreetmap.org Git - nominatim.git/blob - src/nominatim_core/db/connection.py
adapt cmake to new directory structure
[nominatim.git] / src / nominatim_core / db / connection.py
1 # SPDX-License-Identifier: GPL-3.0-or-later
2 #
3 # This file is part of Nominatim. (https://nominatim.org)
4 #
5 # Copyright (C) 2024 by the Nominatim developer community.
6 # For a full list of authors see the git log.
7 """
8 Specialised connection and cursor functions.
9 """
10 from typing import Optional, Any, Callable, ContextManager, Dict, cast, overload, Tuple, Iterable
11 import contextlib
12 import logging
13 import os
14
15 import psycopg2
16 import psycopg2.extensions
17 import psycopg2.extras
18 from psycopg2 import sql as pysql
19
20 from ..typing import SysEnv, Query, T_cursor
21 from ..errors import UsageError
22
23 LOG = logging.getLogger()
24
25 class Cursor(psycopg2.extras.DictCursor):
26     """ A cursor returning dict-like objects and providing specialised
27         execution functions.
28     """
29     # pylint: disable=arguments-renamed,arguments-differ
30     def execute(self, query: Query, args: Any = None) -> None:
31         """ Query execution that logs the SQL query when debugging is enabled.
32         """
33         if LOG.isEnabledFor(logging.DEBUG):
34             LOG.debug(self.mogrify(query, args).decode('utf-8'))
35
36         super().execute(query, args)
37
38
39     def execute_values(self, sql: Query, argslist: Iterable[Tuple[Any, ...]],
40                        template: Optional[Query] = None) -> None:
41         """ Wrapper for the psycopg2 convenience function to execute
42             SQL for a list of values.
43         """
44         LOG.debug("SQL execute_values(%s, %s)", sql, argslist)
45
46         psycopg2.extras.execute_values(self, sql, argslist, template=template)
47
48
49     def scalar(self, sql: Query, args: Any = None) -> Any:
50         """ Execute query that returns a single value. The value is returned.
51             If the query yields more than one row, a ValueError is raised.
52         """
53         self.execute(sql, args)
54
55         if self.rowcount != 1:
56             raise RuntimeError("Query did not return a single row.")
57
58         result = self.fetchone()
59         assert result is not None
60
61         return result[0]
62
63
64     def drop_table(self, name: str, if_exists: bool = True, cascade: bool = False) -> None:
65         """ Drop the table with the given name.
66             Set `if_exists` to False if a non-existent table should raise
67             an exception instead of just being ignored. If 'cascade' is set
68             to True then all dependent tables are deleted as well.
69         """
70         sql = 'DROP TABLE '
71         if if_exists:
72             sql += 'IF EXISTS '
73         sql += '{}'
74         if cascade:
75             sql += ' CASCADE'
76
77         self.execute(pysql.SQL(sql).format(pysql.Identifier(name)))
78
79
80 class Connection(psycopg2.extensions.connection):
81     """ A connection that provides the specialised cursor by default and
82         adds convenience functions for administrating the database.
83     """
84     @overload # type: ignore[override]
85     def cursor(self) -> Cursor:
86         ...
87
88     @overload
89     def cursor(self, name: str) -> Cursor:
90         ...
91
92     @overload
93     def cursor(self, cursor_factory: Callable[..., T_cursor]) -> T_cursor:
94         ...
95
96     def cursor(self, cursor_factory  = Cursor, **kwargs): # type: ignore
97         """ Return a new cursor. By default the specialised cursor is returned.
98         """
99         return super().cursor(cursor_factory=cursor_factory, **kwargs)
100
101
102     def table_exists(self, table: str) -> bool:
103         """ Check that a table with the given name exists in the database.
104         """
105         with self.cursor() as cur:
106             num = cur.scalar("""SELECT count(*) FROM pg_tables
107                                 WHERE tablename = %s and schemaname = 'public'""", (table, ))
108             return num == 1 if isinstance(num, int) else False
109
110
111     def table_has_column(self, table: str, column: str) -> bool:
112         """ Check if the table 'table' exists and has a column with name 'column'.
113         """
114         with self.cursor() as cur:
115             has_column = cur.scalar("""SELECT count(*) FROM information_schema.columns
116                                        WHERE table_name = %s
117                                              and column_name = %s""",
118                                     (table, column))
119             return has_column > 0 if isinstance(has_column, int) else False
120
121
122     def index_exists(self, index: str, table: Optional[str] = None) -> bool:
123         """ Check that an index with the given name exists in the database.
124             If table is not None then the index must relate to the given
125             table.
126         """
127         with self.cursor() as cur:
128             cur.execute("""SELECT tablename FROM pg_indexes
129                            WHERE indexname = %s and schemaname = 'public'""", (index, ))
130             if cur.rowcount == 0:
131                 return False
132
133             if table is not None:
134                 row = cur.fetchone()
135                 if row is None or not isinstance(row[0], str):
136                     return False
137                 return row[0] == table
138
139         return True
140
141
142     def drop_table(self, name: str, if_exists: bool = True, cascade: bool = False) -> None:
143         """ Drop the table with the given name.
144             Set `if_exists` to False if a non-existent table should raise
145             an exception instead of just being ignored.
146         """
147         with self.cursor() as cur:
148             cur.drop_table(name, if_exists, cascade)
149         self.commit()
150
151
152     def server_version_tuple(self) -> Tuple[int, int]:
153         """ Return the server version as a tuple of (major, minor).
154             Converts correctly for pre-10 and post-10 PostgreSQL versions.
155         """
156         version = self.server_version
157         if version < 100000:
158             return (int(version / 10000), int((version % 10000) / 100))
159
160         return (int(version / 10000), version % 10000)
161
162
163     def postgis_version_tuple(self) -> Tuple[int, int]:
164         """ Return the postgis version installed in the database as a
165             tuple of (major, minor). Assumes that the PostGIS extension
166             has been installed already.
167         """
168         with self.cursor() as cur:
169             version = cur.scalar('SELECT postgis_lib_version()')
170
171         version_parts = version.split('.')
172         if len(version_parts) < 2:
173             raise UsageError(f"Error fetching Postgis version. Bad format: {version}")
174
175         return (int(version_parts[0]), int(version_parts[1]))
176
177
178     def extension_loaded(self, extension_name: str) -> bool:
179         """ Return True if the hstore extension is loaded in the database.
180         """
181         with self.cursor() as cur:
182             cur.execute('SELECT extname FROM pg_extension WHERE extname = %s', (extension_name, ))
183             return cur.rowcount > 0
184
185
186 class ConnectionContext(ContextManager[Connection]):
187     """ Context manager of the connection that also provides direct access
188         to the underlying connection.
189     """
190     connection: Connection
191
192 def connect(dsn: str) -> ConnectionContext:
193     """ Open a connection to the database using the specialised connection
194         factory. The returned object may be used in conjunction with 'with'.
195         When used outside a context manager, use the `connection` attribute
196         to get the connection.
197     """
198     try:
199         conn = psycopg2.connect(dsn, connection_factory=Connection)
200         ctxmgr = cast(ConnectionContext, contextlib.closing(conn))
201         ctxmgr.connection = conn
202         return ctxmgr
203     except psycopg2.OperationalError as err:
204         raise UsageError(f"Cannot connect to database: {err}") from err
205
206
207 # Translation from PG connection string parameters to PG environment variables.
208 # Derived from https://www.postgresql.org/docs/current/libpq-envars.html.
209 _PG_CONNECTION_STRINGS = {
210     'host': 'PGHOST',
211     'hostaddr': 'PGHOSTADDR',
212     'port': 'PGPORT',
213     'dbname': 'PGDATABASE',
214     'user': 'PGUSER',
215     'password': 'PGPASSWORD',
216     'passfile': 'PGPASSFILE',
217     'channel_binding': 'PGCHANNELBINDING',
218     'service': 'PGSERVICE',
219     'options': 'PGOPTIONS',
220     'application_name': 'PGAPPNAME',
221     'sslmode': 'PGSSLMODE',
222     'requiressl': 'PGREQUIRESSL',
223     'sslcompression': 'PGSSLCOMPRESSION',
224     'sslcert': 'PGSSLCERT',
225     'sslkey': 'PGSSLKEY',
226     'sslrootcert': 'PGSSLROOTCERT',
227     'sslcrl': 'PGSSLCRL',
228     'requirepeer': 'PGREQUIREPEER',
229     'ssl_min_protocol_version': 'PGSSLMINPROTOCOLVERSION',
230     'ssl_max_protocol_version': 'PGSSLMAXPROTOCOLVERSION',
231     'gssencmode': 'PGGSSENCMODE',
232     'krbsrvname': 'PGKRBSRVNAME',
233     'gsslib': 'PGGSSLIB',
234     'connect_timeout': 'PGCONNECT_TIMEOUT',
235     'target_session_attrs': 'PGTARGETSESSIONATTRS',
236 }
237
238
239 def get_pg_env(dsn: str,
240                base_env: Optional[SysEnv] = None) -> Dict[str, str]:
241     """ Return a copy of `base_env` with the environment variables for
242         PostgreSQL set up from the given database connection string.
243         If `base_env` is None, then the OS environment is used as a base
244         environment.
245     """
246     env = dict(base_env if base_env is not None else os.environ)
247
248     for param, value in psycopg2.extensions.parse_dsn(dsn).items():
249         if param in _PG_CONNECTION_STRINGS:
250             env[_PG_CONNECTION_STRINGS[param]] = value
251         else:
252             LOG.error("Unknown connection parameter '%s' ignored.", param)
253
254     return env