codekingpro/portable-devtools

1# dialects/postgresql/provision.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8
9import time
10
11from ... import exc
12from ... import inspect
13from ... import text
14from ...testing import warn_test_suite
15from ...testing.provision import create_db
16from ...testing.provision import drop_all_schema_objects_post_tables
17from ...testing.provision import drop_all_schema_objects_pre_tables
18from ...testing.provision import drop_db
19from ...testing.provision import log
20from ...testing.provision import post_configure_engine
21from ...testing.provision import prepare_for_drop_tables
22from ...testing.provision import set_default_schema_on_connection
23from ...testing.provision import temp_table_keyword_args
24from ...testing.provision import upsert
25
26
27@create_db.for_db("postgresql")
28def _pg_create_db(cfg, eng, ident):
29    template_db = cfg.options.postgresql_templatedb
30
31    with eng.execution_options(isolation_level="AUTOCOMMIT").begin() as conn:
32        if not template_db:
33            template_db = conn.exec_driver_sql(
34                "select current_database()"
35            ).scalar()
36
37        attempt = 0
38        while True:
39            try:
40                conn.exec_driver_sql(
41                    "CREATE DATABASE %s TEMPLATE %s" % (ident, template_db)
42                )
43            except exc.OperationalError as err:
44                attempt += 1
45                if attempt >= 3:
46                    raise
47                if "accessed by other users" in str(err):
48                    log.info(
49                        "Waiting to create %s, URI %r, "
50                        "template DB %s is in use sleeping for .5",
51                        ident,
52                        eng.url,
53                        template_db,
54                    )
55                    time.sleep(0.5)
56            except:
57                raise
58            else:
59                break
60
61
62@drop_db.for_db("postgresql")
63def _pg_drop_db(cfg, eng, ident):
64    with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
65        with conn.begin():
66            conn.execute(
67                text(
68                    "select pg_terminate_backend(pid) from pg_stat_activity "
69                    "where usename=current_user and pid != pg_backend_pid() "
70                    "and datname=:dname"
71                ),
72                dict(dname=ident),
73            )
74            conn.exec_driver_sql("DROP DATABASE %s" % ident)
75
76
77@temp_table_keyword_args.for_db("postgresql")
78def _postgresql_temp_table_keyword_args(cfg, eng):
79    return {"prefixes": ["TEMPORARY"]}
80
81
82@set_default_schema_on_connection.for_db("postgresql")
83def _postgresql_set_default_schema_on_connection(
84    cfg, dbapi_connection, schema_name
85):
86    existing_autocommit = dbapi_connection.autocommit
87    dbapi_connection.autocommit = True
88    cursor = dbapi_connection.cursor()
89    cursor.execute("SET SESSION search_path='%s'" % schema_name)
90    cursor.close()
91    dbapi_connection.autocommit = existing_autocommit
92
93
94@drop_all_schema_objects_pre_tables.for_db("postgresql")
95def drop_all_schema_objects_pre_tables(cfg, eng):
96    with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
97        for xid in conn.exec_driver_sql(
98            "select gid from pg_prepared_xacts"
99        ).scalars():
100            conn.exec_driver_sql("ROLLBACK PREPARED '%s'" % xid)
101
102
103@drop_all_schema_objects_post_tables.for_db("postgresql")
104def drop_all_schema_objects_post_tables(cfg, eng):
105    from sqlalchemy.dialects import postgresql
106
107    inspector = inspect(eng)
108    with eng.begin() as conn:
109        for enum in inspector.get_enums("*"):
110            conn.execute(
111                postgresql.DropEnumType(
112                    postgresql.ENUM(name=enum["name"], schema=enum["schema"])
113                )
114            )
115
116
117@prepare_for_drop_tables.for_db("postgresql")
118def prepare_for_drop_tables(config, connection):
119    """Ensure there are no locks on the current username/database."""
120
121    result = connection.exec_driver_sql(
122        "select pid, state, wait_event_type, query "
123        # "select pg_terminate_backend(pid), state, wait_event_type "
124        "from pg_stat_activity where "
125        "usename=current_user "
126        "and datname=current_database() and state='idle in transaction' "
127        "and pid != pg_backend_pid()"
128    )
129    rows = result.all()  # noqa
130    if rows:
131        warn_test_suite(
132            "PostgreSQL may not be able to DROP tables due to "
133            "idle in transaction: %s"
134            % ("; ".join(row._mapping["query"] for row in rows))
135        )
136
137
138@upsert.for_db("postgresql")
139def _upsert(
140    cfg,
141    table,
142    returning,
143    *,
144    set_lambda=None,
145    sort_by_parameter_order=False,
146    index_elements=None,
147):
148    from sqlalchemy.dialects.postgresql import insert
149
150    stmt = insert(table)
151
152    table_pk = inspect(table).selectable
153
154    if set_lambda:
155        if index_elements is None:
156            index_elements = table_pk.primary_key
157        stmt = stmt.on_conflict_do_update(
158            index_elements=index_elements, set_=set_lambda(stmt.excluded)
159        )
160    else:
161        stmt = stmt.on_conflict_do_nothing()
162
163    stmt = stmt.returning(
164        *returning, sort_by_parameter_order=sort_by_parameter_order
165    )
166    return stmt
167
168
169_extensions = [
170    ("citext", (13,)),
171    ("hstore", (13,)),
172]
173
174
175@post_configure_engine.for_db("postgresql")
176def _create_citext_extension(url, engine, follower_ident):
177    with engine.connect() as conn:
178        for extension, min_version in _extensions:
179            if conn.dialect.server_version_info >= min_version:
180                conn.execute(
181                    text(f"CREATE EXTENSION IF NOT EXISTS {extension}")
182                )
183                conn.commit()
184