psycodict 1.0.0rc1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
psycodict/base.py ADDED
@@ -0,0 +1,1528 @@
1
+ # -*- coding: utf-8 -*-
2
+ """
3
+ The shared plumbing underneath every psycodict object.
4
+
5
+ :class:`PostgresBase` is the common base of the database, table and
6
+ statistics classes; it owns statement execution through ``_execute``
7
+ (logging, slow-query warnings, commit/rollback bookkeeping and
8
+ reconnection) together with helpers for inspecting tables, indexes and
9
+ constraints. The module also defines the layout of the ``meta_*`` tables --
10
+ the column lists, types and creation statements shared by everything that
11
+ reads or writes them -- and the metadata format version (``META_FORMAT``)
12
+ stamped into ``meta_format``.
13
+ """
14
+ import csv
15
+ import logging
16
+ import re
17
+ import sys
18
+ import time
19
+ from collections import defaultdict
20
+
21
+ from psycopg import (
22
+ ClientCursor,
23
+ DatabaseError,
24
+ InterfaceError,
25
+ OperationalError,
26
+ ProgrammingError,
27
+ NotSupportedError,
28
+ DataError,
29
+ )
30
+ from psycopg.sql import SQL, Identifier, Placeholder, Literal, Composable
31
+
32
+ from .encoding import Json
33
+ from .utils import reraise, DelayCommit, QueryLogFilter
34
+
35
+
36
+ # This dictionary is used when creating new tables
37
+ # The value associated to each type is the typlen from the pg_type table
38
+ # Reverse sorting by this typlen improves space efficiency
39
+ # due to postgres' alignment requirements
40
+ number_types = {
41
+ "int2": 2,
42
+ "smallint": 2,
43
+ "smallserial": 2,
44
+ "serial2": 2,
45
+ "int4": 4,
46
+ "int": 4,
47
+ "integer": 4,
48
+ "serial": 4,
49
+ "serial4": 4,
50
+ "int8": 8,
51
+ "bigint": 8,
52
+ "bigserial": 8,
53
+ "serial8": 8,
54
+ "numeric": -1,
55
+ "decimal": -1,
56
+ "float4": 4,
57
+ "real": 4,
58
+ "float8": 8,
59
+ "double precision": 8,
60
+ }
61
+ types_whitelist = {
62
+ "boolean": 1,
63
+ "bool": 1,
64
+ "text": -1,
65
+ "char": 1,
66
+ "character": 1,
67
+ "character varying": -1,
68
+ "varchar": -1,
69
+ "json": -1,
70
+ "jsonb": -1,
71
+ "xml": -1,
72
+ "date": 4,
73
+ "interval": 16,
74
+ "time": 8,
75
+ "time without time zone": 8,
76
+ "time with time zone": 12,
77
+ "timetz": 12,
78
+ "timestamp": 8,
79
+ "timestamp without time zone": 8,
80
+ "timestamp with time zone": 8,
81
+ "timestamptz": 8,
82
+ "bytea": -1,
83
+ "bit": -1,
84
+ "bit varying": -1,
85
+ "varbit": -1,
86
+ "point": 16,
87
+ "line": 24,
88
+ "lseg": 32,
89
+ "path": -1,
90
+ "box": 32,
91
+ "polygon": -1,
92
+ "circle": 24,
93
+ "tsquery": -1,
94
+ "tsvector": -1,
95
+ "txid_snapshot": -1,
96
+ "uuid": 16,
97
+ "cidr": -1,
98
+ "inet": -1,
99
+ "macaddr": 6,
100
+ "money": 8,
101
+ "pg_lsn": 8,
102
+ }
103
+ types_whitelist.update(number_types)
104
+ # add arrays
105
+ for elt in list(types_whitelist):
106
+ types_whitelist[elt + "[]"] = -1
107
+
108
+
109
+ param_types_whitelist = {
110
+ r"^(bit( varying)?|varbit)\s*\([1-9][0-9]*\)$": -1,
111
+ r'(text|(char(acter)?|character varying|varchar(\s*\(1-9][0-9]*\))?))(\s+collate "(c|posix|[a-z][a-z]_[a-z][a-z](\.[a-z0-9-]+)?)")?': -1,
112
+ r"^interval(\s+year|month|day|hour|minute|second|year to month|day to hour|day to minute|day to second|hour to minute|hour to second|minute to second)?(\s*\([0-6]\))?$": 16,
113
+ r"^timestamp\s*\([0-6]\)(\s+with(out)? time zone)?$": 8,
114
+ r"^time\s*\(([0-9]|10)\)(\s+without time zone)?$": 8,
115
+ r"^time\s*\(([0-9]|10)\)\s+with time zone$": 12,
116
+ r"^(numeric|decimal)\s*\([1-9][0-9]*(,\s*(0|[1-9][0-9]*))?\)$": -1,
117
+ }
118
+ param_types_whitelist = {re.compile(s): cost for (s, cost) in param_types_whitelist.items()}
119
+
120
+ ##################################################################
121
+ # meta_* infrastructure #
122
+ ##################################################################
123
+
124
+
125
+ def jsonb_idx(cols, cols_type):
126
+ """
127
+ The positions in ``cols`` whose type is ``jsonb``, as a tuple of
128
+ indexes. Used to decide which values need json decoding when reading
129
+ rows of the ``meta_*`` tables.
130
+
131
+ INPUT:
132
+
133
+ - ``cols`` -- a list of column names
134
+ - ``cols_type`` -- a dictionary mapping column names to their types
135
+ """
136
+ return tuple(i for i, elt in enumerate(cols) if cols_type[elt] == "jsonb")
137
+
138
+
139
+ # The version of the metadata format described by the constants below: the
140
+ # layout of the meta_* tables, versioned by a single integer aligned with
141
+ # psycodict's major version (format N is introduced by psycodict N.0). The
142
+ # format of a database is stamped into the single-row meta_format table as
143
+ # (version, min_compat), and every connection checks it: an older but
144
+ # compatible format connects with a warning and reduced functionality, while
145
+ # a layout this psycodict cannot safely use is refused. The policy, and the
146
+ # checklist to follow when changing the format, live in MetadataFormats.md.
147
+ #
148
+ # History:
149
+ # 0 -- the baseline (psycodict 0.x): meta_tables, meta_indexes,
150
+ # meta_constraints and their _hist counterparts, with no format stamp.
151
+ # An unstamped database that has meta tables is format 0.
152
+ # 1 -- (psycodict 1.0) meta_indexes/meta_indexes_hist gained a nullable
153
+ # ``whereclause`` column, holding the predicate of a partial index
154
+ # (NULL for an ordinary index). Compatible: against a format-0
155
+ # database everything keeps working except creating partial indexes.
156
+ # Migrate with ``PostgresDatabase.upgrade_metadata`` (or connect with
157
+ # upgrade=True).
158
+ META_FORMAT = 1
159
+
160
+
161
+ _meta_tables_cols = (
162
+ "name",
163
+ "sort",
164
+ "count_cutoff",
165
+ "id_ordered",
166
+ "out_of_order",
167
+ "stats_valid",
168
+ "label_col",
169
+ "total",
170
+ "important",
171
+ "include_nones",
172
+ )
173
+ _meta_tables_cols_notrequired = (
174
+ "count_cutoff",
175
+ "stats_valid",
176
+ "total",
177
+ "important",
178
+ "include_nones",
179
+ )
180
+ # SQL literals giving the default values for the columns above
181
+ _meta_tables_defaults = {
182
+ "count_cutoff": "1000",
183
+ "stats_valid": "true",
184
+ "total": "0",
185
+ "important": "false",
186
+ "include_nones": "true",
187
+ }
188
+ _meta_tables_types = dict(zip(_meta_tables_cols, (
189
+ "text",
190
+ "jsonb",
191
+ "smallint",
192
+ "boolean",
193
+ "boolean",
194
+ "boolean",
195
+ "text",
196
+ "bigint",
197
+ "boolean",
198
+ "boolean",
199
+ )))
200
+ _meta_tables_jsonb_idx = jsonb_idx(_meta_tables_cols, _meta_tables_types)
201
+
202
+ _meta_indexes_cols = (
203
+ "index_name",
204
+ "table_name",
205
+ "type",
206
+ "columns",
207
+ "modifiers",
208
+ "storage_params",
209
+ # The predicate of a partial index (raw SQL), or NULL for an ordinary
210
+ # index. Added in metadata format 1; see META_FORMAT.
211
+ "whereclause",
212
+ )
213
+ _meta_indexes_types = dict(
214
+ zip(_meta_indexes_cols, ("text", "text", "text", "jsonb", "jsonb", "jsonb", "text"))
215
+ )
216
+ _meta_indexes_jsonb_idx = jsonb_idx(_meta_indexes_cols, _meta_indexes_types)
217
+
218
+ _meta_constraints_cols = (
219
+ "constraint_name",
220
+ "table_name",
221
+ "type",
222
+ "columns",
223
+ "check_func",
224
+ )
225
+ _meta_constraints_types = dict(
226
+ zip(_meta_constraints_cols, ("text", "text", "text", "jsonb", "text"))
227
+ )
228
+ _meta_constraints_jsonb_idx = jsonb_idx(_meta_constraints_cols, _meta_constraints_types)
229
+
230
+ # Columns introduced by a metadata format bump: column -> the format that
231
+ # added it; columns not listed are part of the format-0 baseline. A format
232
+ # bump must append its columns at the end of the _cols tuple above (see
233
+ # MetadataFormats.md), so that the columns of an older format are a prefix of
234
+ # the current ones.
235
+ _meta_col_formats = {
236
+ "meta_tables": {},
237
+ "meta_indexes": {"whereclause": 1},
238
+ "meta_constraints": {},
239
+ }
240
+
241
+
242
+ def _meta_cols_types_jsonb_idx(meta_name, fmt=None):
243
+ """
244
+ The (columns, types, jsonb column indexes) of a metadata table.
245
+
246
+ ``fmt`` restricts the columns to those present in that metadata format
247
+ (a prefix of the current ones, since format bumps only append columns);
248
+ the default is the current format. Callers touching a live database
249
+ should pass the connection's format, ``self._db._meta_format``, so that
250
+ their SQL matches the columns the database actually has.
251
+ """
252
+ assert meta_name in ["meta_tables", "meta_indexes", "meta_constraints"]
253
+ if meta_name == "meta_tables":
254
+ meta_cols = _meta_tables_cols
255
+ meta_types = _meta_tables_types
256
+ meta_jsonb_idx = _meta_tables_jsonb_idx
257
+ elif meta_name == "meta_indexes":
258
+ meta_cols = _meta_indexes_cols
259
+ meta_types = _meta_indexes_types
260
+ meta_jsonb_idx = _meta_indexes_jsonb_idx
261
+ elif meta_name == "meta_constraints":
262
+ meta_cols = _meta_constraints_cols
263
+ meta_types = _meta_constraints_types
264
+ meta_jsonb_idx = _meta_constraints_jsonb_idx
265
+
266
+ if fmt is not None and fmt < META_FORMAT:
267
+ added = _meta_col_formats[meta_name]
268
+ meta_cols = tuple(col for col in meta_cols if added.get(col, 0) <= fmt)
269
+ meta_jsonb_idx = jsonb_idx(meta_cols, meta_types)
270
+ return meta_cols, meta_types, meta_jsonb_idx
271
+
272
+
273
+ def _meta_table_name(meta_name):
274
+ meta_cols, _, _ = _meta_cols_types_jsonb_idx(meta_name)
275
+ # the column which will match search_table
276
+ table_name = "table_name"
277
+ if "name" in meta_cols:
278
+ table_name = "name"
279
+ return table_name
280
+
281
+
282
+ class PostgresBase():
283
+ """
284
+ A base class for various objects that interact with Postgres.
285
+
286
+ Any class inheriting from this one must provide a connection
287
+ to the postgres database, as well as a name used when creating a logger.
288
+ """
289
+
290
+ def __init__(self, loggername, db):
291
+ # Have to record this object in the db so that we can reset the connection if necessary.
292
+ # This function also sets self.conn
293
+ db._register_object(self)
294
+ self._db = db
295
+
296
+ logging_options = db.config.options["logging"]
297
+ self.slow_cutoff = logging_options["slowcutoff"]
298
+ self._logger = l = logging.getLogger(loggername)
299
+ l.propagate = False
300
+ # we only want 2 handlers
301
+ l.handlers = []
302
+ l.setLevel(logging_options.get('loglevel', logging.INFO))
303
+ formatter = logging.Formatter("%(asctime)s - %(message)s")
304
+ fhandler = logging.FileHandler(logging_options["slowlogfile"])
305
+ fhandler.setFormatter(formatter)
306
+ fhandler.addFilter(QueryLogFilter())
307
+ l.addHandler(fhandler)
308
+ shandler = logging.StreamHandler()
309
+ shandler.setFormatter(formatter)
310
+ l.addHandler(shandler)
311
+
312
+ def _mogrify(self, query, values):
313
+ """
314
+ Render a query with values interpolated, for logging and error messages.
315
+
316
+ psycopg3 only supports client-side interpolation through ClientCursor,
317
+ so we create a temporary one (psycopg2 had mogrify on every cursor).
318
+ """
319
+ return ClientCursor(self.conn).mogrify(query, values)
320
+
321
+ def _execute(
322
+ self,
323
+ query,
324
+ values=None,
325
+ silent=None,
326
+ values_list=False,
327
+ template=None,
328
+ commit=None,
329
+ slow_note=None,
330
+ reissued=False,
331
+ buffered=False
332
+ ):
333
+ """
334
+ Execute an SQL command, properly catching errors and returning the resulting cursor.
335
+
336
+ INPUT:
337
+
338
+ - ``query`` -- an SQL Composable object, the SQL command to execute.
339
+ - ``values`` -- values to substitute for %s in the query. Quoting from the documentation
340
+ for psycopg2 (https://initd.org/psycopg/docs/usage.html#passing-parameters-to-sql-queries):
341
+
342
+ Never, never, NEVER use Python string concatenation (+) or string parameters
343
+ interpolation (%) to pass variables to a SQL query string. Not even at gunpoint.
344
+
345
+ - ``silent`` -- boolean (default None). If True, don't log a warning for a slow query.
346
+ If None, allow DelayCommit contexts to control silencing.
347
+ - ``values_list`` -- boolean (default False). If True, use the ``execute_values`` method,
348
+ designed for inserting multiple values.
349
+ - ``template`` -- string, for use with ``values_list`` to insert constant values:
350
+ for example ``"(%s, %s, 42)"``. See the documentation of ``execute_values``
351
+ for more details.
352
+ - ``commit`` -- boolean (default None). Whether to commit changes on success. The default
353
+ is to commit unless we are currently in a DelayCommit context.
354
+ - ``slow_note`` -- a tuple for generating more useful data for slow query logging.
355
+ - ``reissued`` -- used internally to prevent infinite recursion when attempting to
356
+ reset the connection.
357
+ - ``buffered`` -- whether to create a server side cursor that must be manually
358
+ closed and connection committed (to closed the transaction) after using it,
359
+ this implies ``commit=False``.
360
+
361
+ .. NOTE:
362
+
363
+ If the Postgres connection has been closed, the execute statement will fail.
364
+ We try to recover gracefully by attempting to open a new connection
365
+ and issuing the command again. However, this approach is not prudent if this
366
+ execute statement is one of a chain of statements, which we detect by checking
367
+ whether ``commit == False``. In this case, we will reset the connection but reraise
368
+ the interface error.
369
+
370
+ The upshot is that you should use ``commit=False`` even for the last of a chain of
371
+ execute statements, then explicitly call ``self.conn.commit()`` afterward.
372
+
373
+ OUTPUT:
374
+
375
+ - a cursor object from which the resulting records can be obtained via iteration.
376
+
377
+ This function will also log slow queries.
378
+ """
379
+ if not isinstance(query, Composable):
380
+ raise TypeError("You must use the psycopg.sql module to execute queries")
381
+
382
+ if buffered:
383
+ if commit is None:
384
+ commit = False
385
+ elif commit:
386
+ raise ValueError("buffered and commit are incompatible")
387
+
388
+ try:
389
+ cur = self._db._cursor(buffered=buffered)
390
+
391
+ t = time.time()
392
+ if values_list:
393
+ # This used to use psycopg2's execute_values; with psycopg3
394
+ # we expand the single "VALUES %s" placeholder to a per-row
395
+ # template and rely on executemany, which batches efficiently
396
+ # using pipeline mode.
397
+ if values:
398
+ if template is not None:
399
+ template = template.as_string(self.conn)
400
+ else:
401
+ template = "(" + ",".join(["%s"] * len(values[0])) + ")"
402
+ cur.executemany(query.as_string(self.conn).replace("%s", template, 1), values)
403
+ else:
404
+ try:
405
+ cur.execute(query, values)
406
+ except (OperationalError, ProgrammingError, NotSupportedError, DataError, SyntaxError) as e:
407
+ try:
408
+ context = " happens while executing {}".format(self._mogrify(query, values))
409
+ except Exception:
410
+ context = " happens while executing {} with values {}".format(query, values)
411
+ reraise(type(e), type(e)(str(e) + context), sys.exc_info()[2])
412
+ if silent is False or (silent is None and not self._db._silenced):
413
+ t = time.time() - t
414
+ if t > self.slow_cutoff:
415
+ if values_list:
416
+ query = query.as_string(self.conn).replace("%s", "VALUES_LIST")
417
+ elif values:
418
+ try:
419
+ query = self._mogrify(query, values)
420
+ except Exception:
421
+ # This shouldn't happen since the execution above was successful
422
+ query = query + str(values)
423
+ else:
424
+ query = query.as_string(self.conn)
425
+ if isinstance(query, bytes): # PY3 compatibility
426
+ query = query.decode("utf-8")
427
+ self._logger.info(query + " ran in \033[91m {0!s}s \033[0m".format(t))
428
+ if slow_note is not None:
429
+ self._logger.info(
430
+ "Replicate with db.%s.%s(%s)",
431
+ slow_note[0],
432
+ slow_note[1],
433
+ ", ".join(str(c) for c in slow_note[2:]),
434
+ )
435
+ except (DatabaseError, InterfaceError):
436
+ if self.conn.closed != 0:
437
+ # If reissued, we need to raise since we're recursing.
438
+ if reissued:
439
+ raise
440
+ # Attempt to reset the connection
441
+ self._db.reset_connection()
442
+ if commit or (commit is None and self._db._nocommit_stack == 0):
443
+ return self._execute(
444
+ query,
445
+ values=values,
446
+ silent=silent,
447
+ values_list=values_list,
448
+ template=template,
449
+ commit=commit,
450
+ slow_note=slow_note,
451
+ buffered=buffered,
452
+ reissued=True,
453
+ )
454
+ else:
455
+ raise
456
+ else:
457
+ self.conn.rollback()
458
+ raise
459
+ else:
460
+ if commit or (commit is None and self._db._nocommit_stack == 0):
461
+ self.conn.commit()
462
+ return cur
463
+
464
+ def _table_exists(self, tablename):
465
+ """
466
+ Check whether the specified table exists
467
+
468
+ INPUT:
469
+
470
+ - ``tablename`` -- a string, the name of the table
471
+ """
472
+ cur = self._execute(SQL("SELECT 1 FROM pg_tables where tablename=%s"), [tablename], silent=True)
473
+ return cur.fetchone() is not None
474
+
475
+ def _all_tablenames(self):
476
+ """
477
+ Return all (postgres) table names in the database
478
+ """
479
+ return [rec[0] for rec in self._execute(SQL("SELECT tablename FROM pg_tables ORDER BY tablename"), silent=True)]
480
+
481
+ def _get_locks(self):
482
+ return self._execute(SQL(
483
+ "SELECT t.relname, l.mode, l.pid, age(clock_timestamp(), a.backend_start) "
484
+ "FROM pg_locks l "
485
+ "JOIN pg_stat_all_tables t ON l.relation = t.relid JOIN pg_stat_activity a ON l.pid = a.pid "
486
+ "WHERE l.granted AND t.schemaname <> 'pg_toast'::name AND t.schemaname <> 'pg_catalog'::name"
487
+ ))
488
+
489
+ def _table_locked(self, tablename, types="all"):
490
+ """
491
+ Tests whether a table is locked.
492
+
493
+ INPUT:
494
+
495
+ - tablename -- a string, the name of the table
496
+ - types -- either a string describing the operation being performed
497
+ (which is translated to a list of lock types with which that operation conflicts)
498
+ or a list of lock types.
499
+
500
+ The valid strings are:
501
+
502
+ - 'update'
503
+ - 'delete'
504
+ - 'insert'
505
+ - 'index'
506
+ - 'select'
507
+ - 'all' (includes all locks)
508
+
509
+ The valid lock types to filter on are:
510
+
511
+ - 'AccessShareLock'
512
+ - 'RowShareLock'
513
+ - 'RowExclusiveLock'
514
+ - 'ShareUpdateExclusiveLock'
515
+ - 'ShareLock'
516
+ - 'ShareRowExclusiveLock'
517
+ - 'ExclusiveLock'
518
+ - 'AccessExclusiveLock'
519
+
520
+ OUTPUT:
521
+
522
+ A list of pairs (locktype, pid) where locktype is a string as above,
523
+ and pid is the process id of the postgres transaction holding the lock.
524
+ """
525
+ if isinstance(types, str):
526
+ if types in ["update", "delete", "insert"]:
527
+ types = [
528
+ "ShareLock",
529
+ "ShareRowExclusiveLock",
530
+ "ExclusiveLock",
531
+ "AccessExclusiveLock",
532
+ ]
533
+ elif types == "index":
534
+ types = [
535
+ "RowExclusiveLock",
536
+ "ShareUpdateExclusiveLock",
537
+ "ShareRowExclusiveLock",
538
+ "ExclusiveLock",
539
+ "AccessExclusiveLock",
540
+ ]
541
+ elif types == "select":
542
+ types = [
543
+ "AccessExclusiveLock"
544
+ ]
545
+ elif types != "all":
546
+ raise ValueError("Invalid lock type")
547
+ if types != "all":
548
+ good_types = [
549
+ "AccessShareLock",
550
+ "RowShareLock",
551
+ "RowExclusiveLock",
552
+ "ShareUpdateExclusiveLock",
553
+ "ShareLock",
554
+ "ShareRowExclusiveLock",
555
+ "ExclusiveLock",
556
+ "AccessExclusiveLock",
557
+ ]
558
+ bad_types = [locktype for locktype in types if locktype not in good_types]
559
+ if bad_types:
560
+ raise ValueError("Invalid lock type(s): %s" % (", ".join(bad_types)))
561
+ return [
562
+ (locktype, pid)
563
+ for (name, locktype, pid, t) in self._get_locks()
564
+ if name == tablename and (types == "all" or locktype in types) and pid != self.conn.info.backend_pid
565
+ ]
566
+
567
+ def _index_exists(self, indexname, tablename=None):
568
+ """
569
+ Check whether the specified index exists
570
+
571
+ INPUT:
572
+
573
+ - ``indexname`` -- a string, the name of the index
574
+ - ``tablename`` -- (optional) a string
575
+
576
+ OUTPUT:
577
+
578
+ If ``tablename`` specified, returns a boolean. If not, returns
579
+ ``False`` if there is no index with this name, or the corresponding tablename
580
+ as a string if there is.
581
+ """
582
+ if tablename:
583
+ cur = self._execute(
584
+ SQL("SELECT 1 FROM pg_indexes WHERE indexname = %s AND tablename = %s"),
585
+ [indexname, tablename],
586
+ silent=True,
587
+ )
588
+ return cur.fetchone() is not None
589
+ else:
590
+ cur = self._execute(
591
+ SQL("SELECT tablename FROM pg_indexes WHERE indexname=%s"),
592
+ [indexname],
593
+ silent=True,
594
+ )
595
+ table = cur.fetchone()
596
+ if table is None:
597
+ return False
598
+ else:
599
+ return table[0]
600
+
601
+ def _relation_exists(self, name):
602
+ """
603
+ Check whether the specified relation exists. Relations are indexes or constraints.
604
+
605
+ INPUT:
606
+
607
+ - ``name`` -- a string, the name of the relation
608
+ """
609
+ cur = self._execute(SQL("SELECT 1 FROM pg_class where relname = %s"), [name])
610
+ return cur.fetchone() is not None
611
+
612
+ def _constraint_exists(self, constraintname, tablename=None):
613
+ """
614
+ Check whether the specified constraint exists
615
+
616
+ INPUT:
617
+
618
+ - ``constraintname`` -- a string, the name of the index
619
+ - ``tablename`` -- (optional) a string
620
+
621
+ OUTPUT:
622
+
623
+ If ``tablename`` specified, returns a boolean. If not, returns
624
+ ``False`` if there is no constraint with this name, or the corresponding tablename
625
+ as a string if there is.
626
+ """
627
+ if tablename:
628
+ cur = self._execute(
629
+ SQL(
630
+ "SELECT 1 from information_schema.table_constraints "
631
+ "WHERE table_name=%s and constraint_name=%s"
632
+ ),
633
+ [tablename, constraintname],
634
+ silent=True,
635
+ )
636
+ return cur.fetchone() is not None
637
+ else:
638
+ cur = self._execute(
639
+ SQL(
640
+ "SELECT table_name from information_schema.table_constraints "
641
+ "WHERE constraint_name=%s"
642
+ ),
643
+ [constraintname],
644
+ silent=True,
645
+ )
646
+ table = cur.fetchone()
647
+ if table is None:
648
+ return False
649
+ else:
650
+ return table[0]
651
+
652
+ def _list_indexes(self, tablename):
653
+ """
654
+ Lists built index names on the search table ``tablename``
655
+ """
656
+ cur = self._execute(
657
+ SQL("SELECT indexname FROM pg_indexes WHERE tablename = %s"),
658
+ [tablename],
659
+ silent=True,
660
+ )
661
+ return [elt[0] for elt in cur]
662
+
663
+ def _list_constraints(self, tablename):
664
+ """
665
+ Lists constraint names on the search table ``tablename``
666
+ """
667
+ # if we look into information_schema.table_constraints
668
+ # we also get internal constraints, I'm not sure why
669
+ # Alternatively, we do a triple join to get the right answer
670
+ cur = self._execute(
671
+ SQL(
672
+ "SELECT con.conname "
673
+ "FROM pg_catalog.pg_constraint con "
674
+ "INNER JOIN pg_catalog.pg_class rel "
675
+ " ON rel.oid = con.conrelid "
676
+ "INNER JOIN pg_catalog.pg_namespace nsp "
677
+ " ON nsp.oid = connamespace "
678
+ "WHERE rel.relname = %s"
679
+ ),
680
+ [tablename],
681
+ silent=True,
682
+ )
683
+ return [elt[0] for elt in cur]
684
+
685
+ def _rename_if_exists(self, name, suffix=""):
686
+ """
687
+ Rename an index or constraint if it exists, appending ``_depN`` if so.
688
+
689
+ INPUT:
690
+
691
+ - ``name`` -- a string, the name of an index or constraint
692
+ - ``suffix`` -- a suffix to append to the name
693
+ """
694
+ if self._relation_exists(name + suffix):
695
+ # First we determine its type
696
+ kind = None
697
+ tablename = self._constraint_exists(name + suffix)
698
+ if tablename:
699
+ kind = "Constraint"
700
+ begin_renamer = SQL("ALTER TABLE {0} RENAME CONSTRAINT").format(Identifier(tablename))
701
+ end_renamer = SQL("{0} TO {1}")
702
+ begin_command = SQL("ALTER TABLE {0}").format(Identifier(tablename))
703
+ end_command = SQL("DROP CONSTRAINT {0}")
704
+ elif self._index_exists(name + suffix):
705
+ kind = "Index"
706
+ begin_renamer = SQL("")
707
+ end_renamer = SQL("ALTER INDEX {0} RENAME TO {1}")
708
+ begin_command = SQL("")
709
+ end_command = SQL("DROP INDEX {0}")
710
+ else:
711
+ raise ValueError(
712
+ "Relation with name "
713
+ + name + suffix
714
+ + " already exists. And it is not an index or a constraint"
715
+ )
716
+
717
+ # Find a new name for the existing index
718
+ depsuffix = "_dep0" + suffix
719
+ i = 0
720
+ deprecated_name = name[: 64 - len(depsuffix)] + depsuffix
721
+ while self._relation_exists(deprecated_name):
722
+ i += 1
723
+ depsuffix = "_dep" + str(i) + suffix
724
+ deprecated_name = name[: 64 - len(depsuffix)] + depsuffix
725
+
726
+ self._execute(
727
+ begin_renamer + end_renamer.format(Identifier(name + suffix), Identifier(deprecated_name))
728
+ )
729
+
730
+ command = begin_command + end_command.format(Identifier(deprecated_name))
731
+
732
+ logging.warning(
733
+ "{} with name {} ".format(kind, name + suffix)
734
+ + "already exists. "
735
+ + "It has been renamed to {} ".format(deprecated_name)
736
+ + "and it can be deleted with the following SQL command:\n"
737
+ + command.as_string(self.conn)
738
+ )
739
+
740
+ def _check_restricted_suffix(self, name, kind="Index", skip_dep=False):
741
+ """
742
+ Checks to ensure that the given name doesn't end with one
743
+ of the following restricted suffixes:
744
+
745
+ - ``_tmp``
746
+ - ``_pkey``
747
+ - ``_oldN``
748
+ - ``_depN``
749
+
750
+ INPUT:
751
+
752
+ - ``name`` -- string, the name of an index or constraint
753
+ - ``kind`` -- either ``"Index"`` or ``"Constraint"`` (only used for error msg)
754
+ - ``skip_dep`` -- if true, allow ``_depN`` as a suffix
755
+ """
756
+ tests = [(r"_old[\d]+$", "_oldN"), (r"_tmp$", "_tmp"), ("_pkey$", "_pkey")]
757
+ if not skip_dep:
758
+ # _rename_if_exists appends "_dep<N>" (no trailing underscore), so
759
+ # the guard must be anchored the same way as its _oldN sibling; the
760
+ # stray trailing "_" here meant it never matched a real deprecated
761
+ # name and the check was dead.
762
+ tests.append((r"_dep[\d]+$", "_depN"))
763
+ for match, message in tests:
764
+ # re.search, not re.match: these patterns are $-anchored
765
+ # suffixes, and match() would only ever find them at the start
766
+ # of the name, so the guard never fired.
767
+ if re.search(match, name):
768
+ raise ValueError(
769
+ "{} name {} is invalid, ".format(kind, name)
770
+ + "cannot end in {}, ".format(message)
771
+ + "try specifying a different name"
772
+ )
773
+
774
+ @staticmethod
775
+ def _sort_str(sort_list):
776
+ """
777
+ Constructs a psycopg.sql.Composable object describing a sort order
778
+ for Postgres from a list of columns.
779
+
780
+ INPUT:
781
+
782
+ - ``sort_list`` -- a list, either of strings (which are interpreted as
783
+ column names in the ascending direction) or of pairs (column name, 1 or -1).
784
+
785
+ OUTPUT:
786
+
787
+ - a Composable to be used by psycopg in the ORDER BY clause.
788
+ """
789
+ PostgresBase._check_sort_duplicates(sort_list)
790
+ L = []
791
+ for col in sort_list:
792
+ if isinstance(col, str):
793
+ L.append(Identifier(col))
794
+ elif col[1] == 1:
795
+ L.append(Identifier(col[0]))
796
+ else:
797
+ L.append(SQL("{0} DESC NULLS LAST").format(Identifier(col[0])))
798
+ return SQL(", ").join(L)
799
+
800
+ @staticmethod
801
+ def _check_sort_duplicates(sort_list):
802
+ """
803
+ Raise if a column appears more than once in ``sort_list`` (a list of
804
+ column names or (column, direction) pairs). A column already fixes the
805
+ order by its first appearance, so a repeat is dead weight and almost
806
+ always a mistake.
807
+ """
808
+ seen = set()
809
+ for col in sort_list:
810
+ name = col if isinstance(col, str) else col[0]
811
+ if name in seen:
812
+ raise ValueError("Duplicate column %r in sort order" % (name,))
813
+ seen.add(name)
814
+
815
+ def _column_types(self, table_name, data_types=None):
816
+ """
817
+ Returns the
818
+ - column list,
819
+ - column types (as a dict), and
820
+ - has_id for a given table_name or list of table names
821
+
822
+ INPUT:
823
+
824
+ - ``table_name`` -- a string or list of strings
825
+ - ``data_types`` -- (optional) a dictionary providing a list of column names and
826
+ types for each table name. If not provided, will be looked up from the database.
827
+
828
+ EXAMPLES::
829
+
830
+ >>> db._column_types('nonexistent')
831
+ ([], {}, False)
832
+ >>> db._column_types('test_fields')
833
+ (['class_group', 'class_number', 'degree', 'disc_abs',
834
+ 'disc_sign', 'label', 'r2', 'ramps'],
835
+ {'id': 'bigint',
836
+ 'class_number': 'integer',
837
+ 'disc_abs': 'integer',
838
+ 'degree': 'smallint',
839
+ 'disc_sign': 'smallint',
840
+ 'r2': 'smallint',
841
+ 'ramps': 'integer[]',
842
+ 'class_group': 'jsonb',
843
+ 'label': 'text'},
844
+ True)
845
+ """
846
+ has_id = False
847
+ col_list = []
848
+ col_type = {}
849
+ if isinstance(table_name, str):
850
+ table_name = [table_name]
851
+ for tname in table_name:
852
+ if data_types is None or tname not in data_types:
853
+ # in case of an array data type, data_type only gives 'ARRAY', while 'udt_name::regtype' gives us 'base_type[]'
854
+ cur = self._execute(
855
+ SQL(
856
+ "SELECT column_name, udt_name::regtype FROM information_schema.columns "
857
+ "WHERE table_name = %s ORDER BY ordinal_position"
858
+ ),
859
+ [tname],
860
+ )
861
+ else:
862
+ cur = data_types[tname]
863
+ for rec in cur:
864
+ col = rec[0]
865
+ if col in col_type and col_type[col] != rec[1]:
866
+ raise ValueError("Type mismatch on %s: %s vs %s" % (col, col_type[col], rec[1]))
867
+ col_type[col] = rec[1]
868
+ if col != "id":
869
+ col_list.append(col)
870
+ else:
871
+ has_id = True
872
+ return sorted(col_list), col_type, has_id
873
+
874
+ def _copy_to_select(self, select, filename, header="", sep="|", silent=False):
875
+ """
876
+ Using COPY ... TO STDOUT, exports the data from a select statement.
877
+
878
+ INPUT:
879
+
880
+ - ``select`` -- an SQL Composable object giving a select statement
881
+ - ``header`` -- An initial header to write to the file
882
+ - ``sep`` -- a separator, defaults to ``|``
883
+ - ``silent`` -- suppress reporting success
884
+ """
885
+ if sep != "\t":
886
+ sep_clause = SQL(" (DELIMITER {0})").format(Literal(sep))
887
+ else:
888
+ sep_clause = SQL("")
889
+ copyto = SQL("COPY ({0}) TO STDOUT{1}").format(select, sep_clause)
890
+ with open(filename, "w") as F:
891
+ try:
892
+ F.write(header)
893
+ cur = self._db._cursor()
894
+ with cur.copy(copyto) as copy:
895
+ for data in copy:
896
+ F.write(bytes(data).decode())
897
+ except Exception:
898
+ self.conn.rollback()
899
+ raise
900
+ else:
901
+ if not silent:
902
+ print("Created file %s" % filename)
903
+
904
+ def _check_header_lines(
905
+ self, F, table_name, columns_set, sep="|", prohibit_missing=True
906
+ ):
907
+ """
908
+ Reads the header lines from a file (row of column names, row of column
909
+ types, blank line), checking if these names match the columns set and
910
+ the types match the expected types in the table.
911
+ Returns a list of column names present in the header.
912
+
913
+ INPUT:
914
+
915
+ - ``F`` -- an open file handle, at the beginning of the file.
916
+ - ``table_name`` -- the table to compare types against (or a list of tables)
917
+ - ``columns_set`` -- a set of the columns expected in the table.
918
+ - ``sep`` -- a string giving the column separator.
919
+ - ``prohibit_missing`` -- raise an error if not all columns present.
920
+
921
+ OUTPUT:
922
+
923
+ The ordered list of columns. The first entry may be ``"id"`` if the data
924
+ contains an id column.
925
+ """
926
+
927
+ col_list, col_type, _ = self._column_types(table_name)
928
+ columns_set.discard("id")
929
+ if not (columns_set <= set(col_list)):
930
+ raise ValueError("{} is not a subset of {}".format(columns_set, col_list))
931
+ header_cols = self._read_header_lines(F, sep=sep)
932
+ names = [elt[0] for elt in header_cols]
933
+ names_set = set(names)
934
+ if "id" in names_set:
935
+ if names[0] != "id":
936
+ raise ValueError("id must be the first column")
937
+ if header_cols[0][1] not in ["int2", "smallint", "int4", "integer", "int8", "bigint"]:
938
+ raise ValueError("id must be of integeral type")
939
+ names_set.discard("id")
940
+ header_cols = header_cols[1:]
941
+
942
+ missing = columns_set - names_set
943
+ extra = names_set - columns_set
944
+ wrong_type = [
945
+ (name, typ)
946
+ for name, typ in header_cols
947
+ if name in columns_set and col_type[name] != typ
948
+ ]
949
+
950
+ if (missing and prohibit_missing) or extra or wrong_type:
951
+ err = ""
952
+ if missing or extra:
953
+ err += "Invalid header: "
954
+ if missing:
955
+ err += ", ".join(list(missing)) + " (missing)"
956
+ if extra:
957
+ err += ", ".join(list(extra)) + " (extra)"
958
+ if wrong_type:
959
+ if len(wrong_type) > 1:
960
+ err += "Invalid types: "
961
+ else:
962
+ err += "Invalid type: "
963
+ err += ", ".join(
964
+ "%s should be %s instead of %s" % (name, col_type[name], typ)
965
+ for name, typ in wrong_type
966
+ )
967
+ raise ValueError(err)
968
+ return names
969
+
970
+ def _copy_from_stdin(self, F, table, columns=None, sep=None, null=r"\N"):
971
+ """
972
+ Stream an open file object into a table using COPY ... FROM STDIN.
973
+
974
+ This replaces psycopg2's ``cursor.copy_from``, which was removed in
975
+ psycopg3 in favor of an explicit COPY statement. Returns the cursor,
976
+ whose ``rowcount`` gives the number of rows loaded.
977
+
978
+ INPUT:
979
+
980
+ - ``F`` -- an open file object to read from
981
+ - ``table`` -- the name of the table to load into
982
+ - ``columns`` -- the columns present in the file, in order
983
+ (defaults to all columns in table order)
984
+ - ``sep`` -- the column separator (defaults to postgres' text-format
985
+ default, a tab, like psycopg2's copy_from did)
986
+ - ``null`` -- the null marker (the text-format default)
987
+ """
988
+ if columns is None:
989
+ cols = SQL("")
990
+ else:
991
+ cols = SQL(" ({0})").format(SQL(", ").join(map(Identifier, columns)))
992
+ if sep is None:
993
+ options = SQL("")
994
+ else:
995
+ options = SQL(" WITH (DELIMITER {0}, NULL {1})").format(Literal(sep), Literal(null))
996
+ copy_sql = SQL("COPY {0}{1} FROM STDIN{2}").format(Identifier(table), cols, options)
997
+ cur = self._db._cursor()
998
+ with cur.copy(copy_sql) as copy:
999
+ while True:
1000
+ chunk = F.read(1 << 20)
1001
+ if not chunk:
1002
+ break
1003
+ copy.write(chunk)
1004
+ return cur
1005
+
1006
+ def _copy_from(self, filename, table, columns, header, kwds):
1007
+ """
1008
+ Helper function for ``copy_from`` and ``reload``.
1009
+
1010
+ INPUT:
1011
+
1012
+ - ``filename`` -- the filename to load
1013
+ - ``table`` -- the table into which the data should be added
1014
+ - ``columns`` -- a list of columns to load (the file may contain them in
1015
+ a different order, specified by a header row)
1016
+ - ``header`` -- whether the file has header rows ordering the columns.
1017
+ This should be True for search tables, False for counts and stats.
1018
+ - ``kwds`` -- may contain ``sep`` and ``null`` options for the COPY
1019
+ """
1020
+ kwds = dict(kwds) # to not modify the dict kwds, with the pop
1021
+ sep = kwds.pop("sep", "|")
1022
+ null = kwds.pop("null", r"\N")
1023
+ kwds.pop("size", None) # psycopg2 buffer size, no longer meaningful
1024
+ if kwds:
1025
+ raise TypeError("Unsupported copy_from options: %s" % ", ".join(kwds))
1026
+
1027
+ with DelayCommit(self, silence=True):
1028
+ with open(filename) as F:
1029
+ if header:
1030
+ # This consumes the first three lines
1031
+ columns = self._check_header_lines(F, table, set(columns), sep=sep)
1032
+ addid = "id" not in columns
1033
+ else:
1034
+ addid = False
1035
+
1036
+ if addid:
1037
+ # create sequence
1038
+ # The values are inlined as literals: DDL statements
1039
+ # cannot take parameters under psycopg3's server-side
1040
+ # binding (psycopg2 interpolated them client-side).
1041
+ cur_count = self.max_id(table)
1042
+ seq_name = table + "_seq"
1043
+ create_seq = SQL(
1044
+ "CREATE SEQUENCE {0} START WITH {1} MINVALUE {1} CACHE 10000"
1045
+ ).format(Identifier(seq_name), Literal(cur_count + 1))
1046
+ self._execute(create_seq)
1047
+ # edit default value
1048
+ alter_table = SQL(
1049
+ "ALTER TABLE {0} ALTER COLUMN {1} SET DEFAULT nextval({2})"
1050
+ ).format(Identifier(table), Identifier("id"), Literal(seq_name))
1051
+ self._execute(alter_table)
1052
+
1053
+ cur = self._copy_from_stdin(F, table, columns, sep, null=null)
1054
+
1055
+ if addid:
1056
+ alter_table = SQL(
1057
+ "ALTER TABLE {0} ALTER COLUMN {1} DROP DEFAULT"
1058
+ ).format(Identifier(table), Identifier("id"))
1059
+ self._execute(alter_table)
1060
+ drop_seq = SQL("DROP SEQUENCE {0}").format(Identifier(seq_name))
1061
+ self._execute(drop_seq)
1062
+
1063
+ return addid, cur.rowcount
1064
+
1065
+ def _get_tablespace(self):
1066
+ # overridden in table and statstable
1067
+ pass
1068
+
1069
+ def _tablespace_clause(self, tablespace=None):
1070
+ """
1071
+ A clause for use in CREATE statements
1072
+ """
1073
+ if tablespace is None:
1074
+ tablespace = self._get_tablespace()
1075
+ if tablespace is None:
1076
+ return SQL("")
1077
+ else:
1078
+ return SQL(" TABLESPACE {0}").format(Identifier(tablespace))
1079
+
1080
+ def _clone(self, table, tmp_table):
1081
+ """
1082
+ Utility function: creates a table with the same schema as the given one.
1083
+
1084
+ INPUT:
1085
+
1086
+ - ``table`` -- string, the name of an existing table
1087
+ - ``tmp_table`` -- string, the name of the new table to create
1088
+ """
1089
+ if self._table_exists(tmp_table):
1090
+ # remove suffix for display message
1091
+ for suffix in ['_counts', '_stats']:
1092
+ if table.endswith(suffix):
1093
+ table = table[:-len(suffix)]
1094
+ raise ValueError(
1095
+ "Temporary table %s already exists. "
1096
+ "Run db.%s.cleanup_from_reload() if you want to delete it and proceed."
1097
+ % (tmp_table, table)
1098
+ )
1099
+ # A bare LIKE copies only the column names and types; carry over the
1100
+ # per-column STORAGE settings (and COMPRESSION, once the server knows
1101
+ # about it) so that clones -- and hence reload and staged, which swap
1102
+ # a clone into place -- do not silently reset them to the defaults.
1103
+ including = SQL(" INCLUDING STORAGE")
1104
+ version = int(self._execute(
1105
+ SQL("SELECT current_setting('server_version_num')"), silent=True
1106
+ ).fetchone()[0])
1107
+ if version >= 140000:
1108
+ # INCLUDING COMPRESSION appeared in PostgreSQL 14 together with
1109
+ # per-column compression itself
1110
+ including += SQL(" INCLUDING COMPRESSION")
1111
+ creator = SQL("CREATE TABLE {0} (LIKE {1}{2}){3}").format(Identifier(tmp_table), Identifier(table), including, self._tablespace_clause())
1112
+ self._execute(creator)
1113
+
1114
+ def _check_col_datatype(self, typ):
1115
+ if typ.lower() not in types_whitelist:
1116
+ if not any(regexp.match(typ.lower()) for regexp in param_types_whitelist):
1117
+ raise RuntimeError("%s is not a valid type" % (typ))
1118
+
1119
+ def _pairs_to_dict(self, L):
1120
+ """
1121
+ Standardize input format for search_columns
1122
+ """
1123
+ if L is None:
1124
+ return L
1125
+ D = defaultdict(list)
1126
+ for (col, typ) in L:
1127
+ D[typ].append(col)
1128
+ return D
1129
+
1130
+ def _get_type_sortkey(self, typ):
1131
+ """
1132
+ Returns the negated storage cost, together with the type
1133
+ Used to sort columns when creating a table for smaller storage footprint
1134
+ """
1135
+ if typ.lower() in types_whitelist:
1136
+ return -types_whitelist[typ.lower()], typ
1137
+ for regexp, cost in param_types_whitelist.items():
1138
+ if regexp.match(typ.lower()):
1139
+ return -cost, typ
1140
+ raise RuntimeError("%s is not a valid type" % (typ))
1141
+
1142
+ def _order_columns(self, coldict, addid="bigint"):
1143
+ """
1144
+ For space reasons, we sort the columns by type, then alphabetically within each type
1145
+
1146
+ This function returns the correct order of the columns.
1147
+ coldict should be in the format output by _pairs_to_dict.
1148
+ """
1149
+ if addid and not any("id" in vals for vals in coldict.values()):
1150
+ if addid not in coldict: # coldict might be a normal dictionary, not a defaultdict
1151
+ coldict[addid] = []
1152
+ coldict[addid].append("id")
1153
+ allcols = []
1154
+ # Note that _get_typlen checks that the type is valid
1155
+ dictorder = sorted(coldict, key=self._get_type_sortkey)
1156
+ for typ in dictorder:
1157
+ for col in sorted(coldict[typ]):
1158
+ # We have whitelisted the types, so it's okay to use string formatting
1159
+ # to insert them into the SQL command.
1160
+ # This is useful so that we can specify the collation in the type
1161
+ allcols.append(SQL("{0} " + typ).format(Identifier(col)))
1162
+ return allcols
1163
+
1164
+ def _create_table(self, name, columns, addid="bigint", tablespace=None):
1165
+ """
1166
+ Utility function: creates a table with the schema specified by ``columns``.
1167
+
1168
+ If self is a table, the new table will be in the same tablespace.
1169
+
1170
+ INPUT:
1171
+
1172
+ - ``name`` -- the desired name
1173
+ - ``columns`` -- list of pairs, where the first entry is
1174
+ the column name and the second one is the corresponding type
1175
+ """
1176
+ if not isinstance(columns, dict):
1177
+ columns = self._pairs_to_dict(columns)
1178
+ ordered = self._order_columns(columns, addid=addid)
1179
+
1180
+ table_col = SQL(", ").join(self._order_columns(columns, addid=addid))
1181
+ creator = SQL("CREATE TABLE {0} ({1}){2}").format(Identifier(name), table_col, self._tablespace_clause(tablespace))
1182
+ self._execute(creator)
1183
+
1184
+ def _create_table_from_header(self, filename, name, sep, addid="bigint", tablespace=None):
1185
+ """
1186
+ Utility function: creates a table with the schema specified in the header of the file.
1187
+ Returns column names found in the header
1188
+
1189
+ INPUT:
1190
+
1191
+ - ``filename`` -- a string, the filename to load the table from
1192
+ - ``name`` -- the name of the table
1193
+ - ``sep`` -- the separator character, defaulting to tab
1194
+ - ``addid`` -- if true, also adds an id column to the created table with the given type
1195
+
1196
+ OUTPUT:
1197
+
1198
+ The list of column names and types found in the header
1199
+ """
1200
+ if self._table_exists(name):
1201
+ error_msg = "Table %s already exists." % name
1202
+ if name.endswith("_tmp"):
1203
+ error_msg += (
1204
+ "Run db.%s.cleanup_from_reload() "
1205
+ "if you want to delete it and proceed." % (name[:-4])
1206
+ )
1207
+ raise ValueError(error_msg)
1208
+ with open(filename, "r") as F:
1209
+ columns = self._read_header_lines(F, sep)
1210
+ col_list = [elt[0] for elt in columns]
1211
+
1212
+ self._create_table(name, columns, addid=addid, tablespace=tablespace)
1213
+ return col_list
1214
+
1215
+ def _swap(self, tables, source, target):
1216
+ """
1217
+ Renames tables, indexes, constraints and primary keys, for use in reload.
1218
+
1219
+ INPUT:
1220
+
1221
+ - ``tables`` -- a list of table names to reload (including suffixes like
1222
+ ``_extra`` or ``_counts`` but not ``_tmp``).
1223
+ - ``source`` -- the source suffix for the swap.
1224
+ - ``target`` -- the target suffix for the swap.
1225
+ """
1226
+ rename_table = SQL("ALTER TABLE {0} RENAME TO {1}")
1227
+ rename_constraint = SQL("ALTER TABLE {0} RENAME CONSTRAINT {1} TO {2}")
1228
+ rename_index = SQL("ALTER INDEX {0} RENAME TO {1}")
1229
+
1230
+ def target_name(name, tablename, kind):
1231
+ original_name = name[:]
1232
+ if source != "" and name.endswith(source):
1233
+ # drop the suffix
1234
+ original_name = original_name[: -len(source)]
1235
+ assert original_name + source == name
1236
+ elif source != "":
1237
+ logging.warning(
1238
+ "{} of {} with name {}".format(kind, tablename, name)
1239
+ + " does not end with the suffix {}".format(source)
1240
+ )
1241
+
1242
+ target_name = original_name + target
1243
+ try:
1244
+ self._check_restricted_suffix(original_name, kind, skip_dep=True)
1245
+ except ValueError:
1246
+ logging.warning(
1247
+ "{} of {} with name {}".format(kind, tablename, name)
1248
+ + " uses a restricted suffix. "
1249
+ + "The name will be extended with a _ in the swap"
1250
+ )
1251
+ target_name = original_name + "_" + target
1252
+ return target_name
1253
+
1254
+ with DelayCommit(self, silence=True):
1255
+ for table in tables:
1256
+ tablename_old = table + source
1257
+ tablename_new = table + target
1258
+ self._execute(rename_table.format(Identifier(tablename_old), Identifier(tablename_new)))
1259
+
1260
+ done = set() # done constraints/indexes
1261
+ # We threat pkey separately
1262
+ pkey_old = table + source + "_pkey"
1263
+ pkey_new = table + target + "_pkey"
1264
+ if self._constraint_exists(pkey_old, tablename_new):
1265
+ self._execute(
1266
+ rename_constraint.format(
1267
+ Identifier(tablename_new),
1268
+ Identifier(pkey_old),
1269
+ Identifier(pkey_new),
1270
+ )
1271
+ )
1272
+ done.add(pkey_new)
1273
+
1274
+ for constraint in self._list_constraints(tablename_new):
1275
+ if constraint in done:
1276
+ continue
1277
+ c_target = target_name(constraint, tablename_new, "Constraint")
1278
+ if c_target != constraint:
1279
+ self._rename_if_exists(c_target)
1280
+ self._execute(
1281
+ rename_constraint.format(
1282
+ Identifier(tablename_new),
1283
+ Identifier(constraint),
1284
+ Identifier(c_target),
1285
+ )
1286
+ )
1287
+ done.add(c_target)
1288
+
1289
+ for index in self._list_indexes(tablename_new):
1290
+ if index in done:
1291
+ continue
1292
+ i_target = target_name(index, tablename_new, "Index")
1293
+ if i_target != index:
1294
+ self._rename_if_exists(i_target)
1295
+ self._execute(
1296
+ rename_index.format(Identifier(index), Identifier(i_target))
1297
+ )
1298
+ done.add(i_target) # not really needed
1299
+
1300
+ def _read_header_lines(self, F, sep="|"):
1301
+ """
1302
+ Reads the header lines from a file
1303
+ (row of column names, row of column types, blank line).
1304
+ Returning the dictionary of columns and their types.
1305
+
1306
+ INPUT:
1307
+
1308
+ - ``F`` -- an open file handle, at the beginning of the file.
1309
+ - ``sep`` -- a string giving the column separator.
1310
+
1311
+ OUTPUT:
1312
+
1313
+ A list of pairs where the first entry is the column and the second the
1314
+ corresponding type
1315
+ """
1316
+ names = [x.strip() for x in F.readline().strip().split(sep)]
1317
+ types = [x.strip() for x in F.readline().strip().split(sep)]
1318
+ blank = F.readline()
1319
+ if blank.strip():
1320
+ raise ValueError("The third line must be blank")
1321
+ if len(names) != len(types):
1322
+ raise ValueError(
1323
+ "The first line specifies %s columns, while the second specifies %s"
1324
+ % (len(names), len(types))
1325
+ )
1326
+ return list(zip(names, types))
1327
+
1328
+ ##################################################################
1329
+ # Exporting, importing, reloading and reverting meta_* #
1330
+ ##################################################################
1331
+
1332
+ def _copy_to_meta(self, meta_name, filename, search_table, sep="|"):
1333
+ # The columns this database actually has: an export from an
1334
+ # older-format database carries that format's columns (a prefix of
1335
+ # the current ones), which _meta_file_columns recognizes on import.
1336
+ meta_cols, _, _ = _meta_cols_types_jsonb_idx(meta_name, self._db._meta_format)
1337
+ table_name = _meta_table_name(meta_name)
1338
+ table_name_sql = Identifier(table_name)
1339
+ meta_name_sql = Identifier(meta_name)
1340
+ cols_sql = SQL(", ").join(map(Identifier, meta_cols))
1341
+ select = SQL("SELECT {} FROM {} WHERE {} = {}").format(
1342
+ cols_sql, meta_name_sql, table_name_sql, Literal(search_table)
1343
+ )
1344
+ now = time.time()
1345
+ with DelayCommit(self):
1346
+ self._copy_to_select(select, filename, sep=sep, silent=True)
1347
+ print(
1348
+ "Exported %s for %s in %.3f secs"
1349
+ % (meta_name, search_table, time.time() - now)
1350
+ )
1351
+
1352
+ def _meta_file_columns(self, meta_name, filename, sep="|"):
1353
+ """
1354
+ The columns of ``meta_name`` that an exported metadata file carries.
1355
+
1356
+ Metadata files have no header line, so the format they were exported
1357
+ at is recovered from their width: format bumps only append columns,
1358
+ so a file written at format f holds the first ``len(columns at f)``
1359
+ of the current columns. Returns that column prefix, or None for an
1360
+ empty file. A file wider than this database's meta table (exported
1361
+ from a newer format than the database is at) or of a width matching
1362
+ no known format is rejected here, with instructions, rather than
1363
+ passed on to COPY to fail cryptically.
1364
+ """
1365
+ with open(filename) as F:
1366
+ first = next(csv.reader(F, delimiter=str(sep)), None)
1367
+ if first is None:
1368
+ return None
1369
+ width = len(first)
1370
+ db_cols, _, _ = _meta_cols_types_jsonb_idx(meta_name, self._db._meta_format)
1371
+ # width -> the oldest format with that many columns
1372
+ widths = {}
1373
+ for fmt in range(META_FORMAT + 1):
1374
+ widths.setdefault(len(_meta_cols_types_jsonb_idx(meta_name, fmt)[0]), fmt)
1375
+ if width not in widths:
1376
+ raise ValueError(
1377
+ "The file %s has %s columns, which matches no known format of "
1378
+ "%s (expected %s)"
1379
+ % (filename, width, meta_name,
1380
+ " or ".join(str(w) for w in sorted(widths)))
1381
+ )
1382
+ if width > len(db_cols):
1383
+ raise ValueError(
1384
+ "The file %s was exported from a database using metadata "
1385
+ "format %s, but this database uses the older format %s: "
1386
+ "migrate it with upgrade_metadata() (or reconnect with "
1387
+ "upgrade=True) before reloading, or re-export the file from "
1388
+ "a format-%s database."
1389
+ % (filename, widths[width], self._db._meta_format,
1390
+ self._db._meta_format)
1391
+ )
1392
+ return db_cols[:width]
1393
+
1394
+ def _copy_from_meta(self, meta_name, filename, sep="|"):
1395
+ # Take the column list from the file's width, so files exported from
1396
+ # an older metadata format keep loading after the database migrates
1397
+ # (columns the file predates are left NULL).
1398
+ meta_cols = self._meta_file_columns(meta_name, filename, sep)
1399
+ if meta_cols is None:
1400
+ return
1401
+ try:
1402
+ with open(filename) as F:
1403
+ self._copy_from_stdin(F, meta_name, meta_cols, sep)
1404
+ except Exception:
1405
+ self.conn.rollback()
1406
+ raise
1407
+
1408
+ def _get_current_meta_version(self, meta_name, search_table):
1409
+ # the column which will match search_table
1410
+ table_name = _meta_table_name(meta_name)
1411
+ table_name_sql = Identifier(table_name)
1412
+ meta_name_hist_sql = Identifier(meta_name + "_hist")
1413
+ res = self._execute(
1414
+ SQL("SELECT MAX(version) FROM {} WHERE {} = %s").format(
1415
+ meta_name_hist_sql, table_name_sql
1416
+ ),
1417
+ [search_table],
1418
+ ).fetchone()[0]
1419
+ if res is None:
1420
+ res = -1
1421
+ return res
1422
+
1423
+ def _reload_meta(self, meta_name, filename, search_table, sep="|"):
1424
+ # The database's columns for the SELECT/INSERT below; the file may
1425
+ # carry fewer (it was exported from an older format), in which case
1426
+ # the trailing columns load as NULL.
1427
+ meta_cols, _, jsonb_idx = _meta_cols_types_jsonb_idx(meta_name, self._db._meta_format)
1428
+ file_cols = self._meta_file_columns(meta_name, filename, sep)
1429
+ # the column which will match search_table
1430
+ table_name = _meta_table_name(meta_name)
1431
+
1432
+ table_name_idx = meta_cols.index(table_name)
1433
+ table_name_sql = Identifier(table_name)
1434
+ meta_name_sql = Identifier(meta_name)
1435
+ meta_name_hist_sql = Identifier(meta_name + "_hist")
1436
+
1437
+ with open(filename, "r") as F:
1438
+ lines = list(csv.reader(F, delimiter=str(sep)))
1439
+ if not lines:
1440
+ return
1441
+ for line in lines:
1442
+ if line[table_name_idx] != search_table:
1443
+ raise RuntimeError(
1444
+ f"column {table_name_idx} (= {line[table_name_idx]}) "
1445
+ f"in the file {filename} doesn't match "
1446
+ f"the search table name {search_table}"
1447
+ )
1448
+
1449
+ with DelayCommit(self, silence=True):
1450
+ # delete the current columns
1451
+ self._execute(
1452
+ SQL("DELETE FROM {} WHERE {} = %s").format(meta_name_sql, table_name_sql),
1453
+ [search_table],
1454
+ )
1455
+
1456
+ # insert new columns
1457
+ with open(filename, "r") as F:
1458
+ try:
1459
+ self._copy_from_stdin(F, meta_name, file_cols, sep)
1460
+ except Exception:
1461
+ self.conn.rollback()
1462
+ raise
1463
+
1464
+ version = self._get_current_meta_version(meta_name, search_table) + 1
1465
+
1466
+ # copy the new rows to history
1467
+ cols_sql = SQL(", ").join(map(Identifier, meta_cols))
1468
+ rows = self._execute(
1469
+ SQL("SELECT {} FROM {} WHERE {} = %s").format(cols_sql, meta_name_sql, table_name_sql),
1470
+ [search_table],
1471
+ )
1472
+
1473
+ cols = meta_cols + ("version",)
1474
+ cols_sql = SQL(", ").join(map(Identifier, cols))
1475
+ place_holder = SQL(", ").join(Placeholder() * len(cols))
1476
+ query = SQL("INSERT INTO {} ({}) VALUES ({})").format(meta_name_hist_sql, cols_sql, place_holder)
1477
+
1478
+ for row in rows:
1479
+ row = [
1480
+ Json(elt) if i in jsonb_idx else elt for i, elt in enumerate(row)
1481
+ ]
1482
+ self._execute(query, row + [version])
1483
+
1484
+ def _revert_meta(self, meta_name, search_table, version=None):
1485
+ meta_cols, _, jsonb_idx = _meta_cols_types_jsonb_idx(meta_name, self._db._meta_format)
1486
+ # the column which will match search_table
1487
+ table_name = _meta_table_name(meta_name)
1488
+
1489
+ table_name_sql = Identifier(table_name)
1490
+ meta_name_sql = Identifier(meta_name)
1491
+ meta_name_hist_sql = Identifier(meta_name + "_hist")
1492
+
1493
+ # by the default goes back one step
1494
+ currentversion = self._get_current_meta_version(meta_name, search_table)
1495
+ if currentversion == -1:
1496
+ raise RuntimeError("No history to revert")
1497
+ if version is None:
1498
+ version = max(0, currentversion - 1)
1499
+
1500
+ with DelayCommit(self, silence=True):
1501
+ # delete current rows
1502
+ self._execute(
1503
+ SQL("DELETE FROM {} WHERE {} = %s").format(meta_name_sql, table_name_sql),
1504
+ [search_table],
1505
+ )
1506
+
1507
+ # copy data from history
1508
+ cols_sql = SQL(", ").join(map(Identifier, meta_cols))
1509
+ rows = self._execute(
1510
+ SQL("SELECT {} FROM {} WHERE {} = %s AND version = %s").format(
1511
+ cols_sql, meta_name_hist_sql, table_name_sql
1512
+ ),
1513
+ [search_table, version],
1514
+ )
1515
+
1516
+ place_holder = SQL(", ").join(Placeholder() * len(meta_cols))
1517
+ query = SQL("INSERT INTO {} ({}) VALUES ({})").format(meta_name_sql, cols_sql, place_holder)
1518
+
1519
+ cols = meta_cols + ("version",)
1520
+ cols_sql = SQL(", ").join(map(Identifier, cols))
1521
+ place_holder = SQL(", ").join(Placeholder() * len(cols))
1522
+ query_hist = SQL("INSERT INTO {} ({}) VALUES ({})").format(
1523
+ meta_name_hist_sql, cols_sql, place_holder
1524
+ )
1525
+ for row in rows:
1526
+ row = [Json(elt) if i in jsonb_idx else elt for i, elt in enumerate(row)]
1527
+ self._execute(query, row)
1528
+ self._execute(query_hist, row + [currentversion + 1])