SQLAlchemy 2.0.36__cp313-cp313-win32.whl

sqlalchemy/dialects/postgresql/pg_catalog.py

@@ -0,0 +1,300 @@
+# dialects/postgresql/pg_catalog.py
+# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: ignore-errors
+
+from .array import ARRAY
+from .types import OID
+from .types import REGCLASS
+from ... import Column
+from ... import func
+from ... import MetaData
+from ... import Table
+from ...types import BigInteger
+from ...types import Boolean
+from ...types import CHAR
+from ...types import Float
+from ...types import Integer
+from ...types import SmallInteger
+from ...types import String
+from ...types import Text
+from ...types import TypeDecorator
+
+
+# types
+class NAME(TypeDecorator):
+    impl = String(64, collation="C")
+    cache_ok = True
+
+
+class PG_NODE_TREE(TypeDecorator):
+    impl = Text(collation="C")
+    cache_ok = True
+
+
+class INT2VECTOR(TypeDecorator):
+    impl = ARRAY(SmallInteger)
+    cache_ok = True
+
+
+class OIDVECTOR(TypeDecorator):
+    impl = ARRAY(OID)
+    cache_ok = True
+
+
+class _SpaceVector:
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if value is None:
+                return value
+            return [int(p) for p in value.split(" ")]
+
+        return process
+
+
+REGPROC = REGCLASS  # seems an alias
+
+# functions
+_pg_cat = func.pg_catalog
+quote_ident = _pg_cat.quote_ident
+pg_table_is_visible = _pg_cat.pg_table_is_visible
+pg_type_is_visible = _pg_cat.pg_type_is_visible
+pg_get_viewdef = _pg_cat.pg_get_viewdef
+pg_get_serial_sequence = _pg_cat.pg_get_serial_sequence
+format_type = _pg_cat.format_type
+pg_get_expr = _pg_cat.pg_get_expr
+pg_get_constraintdef = _pg_cat.pg_get_constraintdef
+pg_get_indexdef = _pg_cat.pg_get_indexdef
+
+# constants
+RELKINDS_TABLE_NO_FOREIGN = ("r", "p")
+RELKINDS_TABLE = RELKINDS_TABLE_NO_FOREIGN + ("f",)
+RELKINDS_VIEW = ("v",)
+RELKINDS_MAT_VIEW = ("m",)
+RELKINDS_ALL_TABLE_LIKE = RELKINDS_TABLE + RELKINDS_VIEW + RELKINDS_MAT_VIEW
+
+# tables
+pg_catalog_meta = MetaData(schema="pg_catalog")
+
+pg_namespace = Table(
+    "pg_namespace",
+    pg_catalog_meta,
+    Column("oid", OID),
+    Column("nspname", NAME),
+    Column("nspowner", OID),
+)
+
+pg_class = Table(
+    "pg_class",
+    pg_catalog_meta,
+    Column("oid", OID, info={"server_version": (9, 3)}),
+    Column("relname", NAME),
+    Column("relnamespace", OID),
+    Column("reltype", OID),
+    Column("reloftype", OID),
+    Column("relowner", OID),
+    Column("relam", OID),
+    Column("relfilenode", OID),
+    Column("reltablespace", OID),
+    Column("relpages", Integer),
+    Column("reltuples", Float),
+    Column("relallvisible", Integer, info={"server_version": (9, 2)}),
+    Column("reltoastrelid", OID),
+    Column("relhasindex", Boolean),
+    Column("relisshared", Boolean),
+    Column("relpersistence", CHAR, info={"server_version": (9, 1)}),
+    Column("relkind", CHAR),
+    Column("relnatts", SmallInteger),
+    Column("relchecks", SmallInteger),
+    Column("relhasrules", Boolean),
+    Column("relhastriggers", Boolean),
+    Column("relhassubclass", Boolean),
+    Column("relrowsecurity", Boolean),
+    Column("relforcerowsecurity", Boolean, info={"server_version": (9, 5)}),
+    Column("relispopulated", Boolean, info={"server_version": (9, 3)}),
+    Column("relreplident", CHAR, info={"server_version": (9, 4)}),
+    Column("relispartition", Boolean, info={"server_version": (10,)}),
+    Column("relrewrite", OID, info={"server_version": (11,)}),
+    Column("reloptions", ARRAY(Text)),
+)
+
+pg_type = Table(
+    "pg_type",
+    pg_catalog_meta,
+    Column("oid", OID, info={"server_version": (9, 3)}),
+    Column("typname", NAME),
+    Column("typnamespace", OID),
+    Column("typowner", OID),
+    Column("typlen", SmallInteger),
+    Column("typbyval", Boolean),
+    Column("typtype", CHAR),
+    Column("typcategory", CHAR),
+    Column("typispreferred", Boolean),
+    Column("typisdefined", Boolean),
+    Column("typdelim", CHAR),
+    Column("typrelid", OID),
+    Column("typelem", OID),
+    Column("typarray", OID),
+    Column("typinput", REGPROC),
+    Column("typoutput", REGPROC),
+    Column("typreceive", REGPROC),
+    Column("typsend", REGPROC),
+    Column("typmodin", REGPROC),
+    Column("typmodout", REGPROC),
+    Column("typanalyze", REGPROC),
+    Column("typalign", CHAR),
+    Column("typstorage", CHAR),
+    Column("typnotnull", Boolean),
+    Column("typbasetype", OID),
+    Column("typtypmod", Integer),
+    Column("typndims", Integer),
+    Column("typcollation", OID, info={"server_version": (9, 1)}),
+    Column("typdefault", Text),
+)
+
+pg_index = Table(
+    "pg_index",
+    pg_catalog_meta,
+    Column("indexrelid", OID),
+    Column("indrelid", OID),
+    Column("indnatts", SmallInteger),
+    Column("indnkeyatts", SmallInteger, info={"server_version": (11,)}),
+    Column("indisunique", Boolean),
+    Column("indnullsnotdistinct", Boolean, info={"server_version": (15,)}),
+    Column("indisprimary", Boolean),
+    Column("indisexclusion", Boolean, info={"server_version": (9, 1)}),
+    Column("indimmediate", Boolean),
+    Column("indisclustered", Boolean),
+    Column("indisvalid", Boolean),
+    Column("indcheckxmin", Boolean),
+    Column("indisready", Boolean),
+    Column("indislive", Boolean, info={"server_version": (9, 3)}),  # 9.3
+    Column("indisreplident", Boolean),
+    Column("indkey", INT2VECTOR),
+    Column("indcollation", OIDVECTOR, info={"server_version": (9, 1)}),  # 9.1
+    Column("indclass", OIDVECTOR),
+    Column("indoption", INT2VECTOR),
+    Column("indexprs", PG_NODE_TREE),
+    Column("indpred", PG_NODE_TREE),
+)
+
+pg_attribute = Table(
+    "pg_attribute",
+    pg_catalog_meta,
+    Column("attrelid", OID),
+    Column("attname", NAME),
+    Column("atttypid", OID),
+    Column("attstattarget", Integer),
+    Column("attlen", SmallInteger),
+    Column("attnum", SmallInteger),
+    Column("attndims", Integer),
+    Column("attcacheoff", Integer),
+    Column("atttypmod", Integer),
+    Column("attbyval", Boolean),
+    Column("attstorage", CHAR),
+    Column("attalign", CHAR),
+    Column("attnotnull", Boolean),
+    Column("atthasdef", Boolean),
+    Column("atthasmissing", Boolean, info={"server_version": (11,)}),
+    Column("attidentity", CHAR, info={"server_version": (10,)}),
+    Column("attgenerated", CHAR, info={"server_version": (12,)}),
+    Column("attisdropped", Boolean),
+    Column("attislocal", Boolean),
+    Column("attinhcount", Integer),
+    Column("attcollation", OID, info={"server_version": (9, 1)}),
+)
+
+pg_constraint = Table(
+    "pg_constraint",
+    pg_catalog_meta,
+    Column("oid", OID),  # 9.3
+    Column("conname", NAME),
+    Column("connamespace", OID),
+    Column("contype", CHAR),
+    Column("condeferrable", Boolean),
+    Column("condeferred", Boolean),
+    Column("convalidated", Boolean, info={"server_version": (9, 1)}),
+    Column("conrelid", OID),
+    Column("contypid", OID),
+    Column("conindid", OID),
+    Column("conparentid", OID, info={"server_version": (11,)}),
+    Column("confrelid", OID),
+    Column("confupdtype", CHAR),
+    Column("confdeltype", CHAR),
+    Column("confmatchtype", CHAR),
+    Column("conislocal", Boolean),
+    Column("coninhcount", Integer),
+    Column("connoinherit", Boolean, info={"server_version": (9, 2)}),
+    Column("conkey", ARRAY(SmallInteger)),
+    Column("confkey", ARRAY(SmallInteger)),
+)
+
+pg_sequence = Table(
+    "pg_sequence",
+    pg_catalog_meta,
+    Column("seqrelid", OID),
+    Column("seqtypid", OID),
+    Column("seqstart", BigInteger),
+    Column("seqincrement", BigInteger),
+    Column("seqmax", BigInteger),
+    Column("seqmin", BigInteger),
+    Column("seqcache", BigInteger),
+    Column("seqcycle", Boolean),
+    info={"server_version": (10,)},
+)
+
+pg_attrdef = Table(
+    "pg_attrdef",
+    pg_catalog_meta,
+    Column("oid", OID, info={"server_version": (9, 3)}),
+    Column("adrelid", OID),
+    Column("adnum", SmallInteger),
+    Column("adbin", PG_NODE_TREE),
+)
+
+pg_description = Table(
+    "pg_description",
+    pg_catalog_meta,
+    Column("objoid", OID),
+    Column("classoid", OID),
+    Column("objsubid", Integer),
+    Column("description", Text(collation="C")),
+)
+
+pg_enum = Table(
+    "pg_enum",
+    pg_catalog_meta,
+    Column("oid", OID, info={"server_version": (9, 3)}),
+    Column("enumtypid", OID),
+    Column("enumsortorder", Float(), info={"server_version": (9, 1)}),
+    Column("enumlabel", NAME),
+)
+
+pg_am = Table(
+    "pg_am",
+    pg_catalog_meta,
+    Column("oid", OID, info={"server_version": (9, 3)}),
+    Column("amname", NAME),
+    Column("amhandler", REGPROC, info={"server_version": (9, 6)}),
+    Column("amtype", CHAR, info={"server_version": (9, 6)}),
+)
+
+pg_collation = Table(
+    "pg_collation",
+    pg_catalog_meta,
+    Column("oid", OID, info={"server_version": (9, 3)}),
+    Column("collname", NAME),
+    Column("collnamespace", OID),
+    Column("collowner", OID),
+    Column("collprovider", CHAR, info={"server_version": (10,)}),
+    Column("collisdeterministic", Boolean, info={"server_version": (12,)}),
+    Column("collencoding", Integer),
+    Column("collcollate", Text),
+    Column("collctype", Text),
+    Column("colliculocale", Text),
+    Column("collicurules", Text, info={"server_version": (16,)}),
+    Column("collversion", Text, info={"server_version": (10,)}),
+)

sqlalchemy/dialects/postgresql/provision.py

@@ -0,0 +1,175 @@
+# dialects/postgresql/provision.py
+# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: ignore-errors
+
+import time
+
+from ... import exc
+from ... import inspect
+from ... import text
+from ...testing import warn_test_suite
+from ...testing.provision import create_db
+from ...testing.provision import drop_all_schema_objects_post_tables
+from ...testing.provision import drop_all_schema_objects_pre_tables
+from ...testing.provision import drop_db
+from ...testing.provision import log
+from ...testing.provision import post_configure_engine
+from ...testing.provision import prepare_for_drop_tables
+from ...testing.provision import set_default_schema_on_connection
+from ...testing.provision import temp_table_keyword_args
+from ...testing.provision import upsert
+
+
+@create_db.for_db("postgresql")
+def _pg_create_db(cfg, eng, ident):
+    template_db = cfg.options.postgresql_templatedb
+
+    with eng.execution_options(isolation_level="AUTOCOMMIT").begin() as conn:
+        if not template_db:
+            template_db = conn.exec_driver_sql(
+                "select current_database()"
+            ).scalar()
+
+        attempt = 0
+        while True:
+            try:
+                conn.exec_driver_sql(
+                    "CREATE DATABASE %s TEMPLATE %s" % (ident, template_db)
+                )
+            except exc.OperationalError as err:
+                attempt += 1
+                if attempt >= 3:
+                    raise
+                if "accessed by other users" in str(err):
+                    log.info(
+                        "Waiting to create %s, URI %r, "
+                        "template DB %s is in use sleeping for .5",
+                        ident,
+                        eng.url,
+                        template_db,
+                    )
+                    time.sleep(0.5)
+            except:
+                raise
+            else:
+                break
+
+
+@drop_db.for_db("postgresql")
+def _pg_drop_db(cfg, eng, ident):
+    with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
+        with conn.begin():
+            conn.execute(
+                text(
+                    "select pg_terminate_backend(pid) from pg_stat_activity "
+                    "where usename=current_user and pid != pg_backend_pid() "
+                    "and datname=:dname"
+                ),
+                dict(dname=ident),
+            )
+            conn.exec_driver_sql("DROP DATABASE %s" % ident)
+
+
+@temp_table_keyword_args.for_db("postgresql")
+def _postgresql_temp_table_keyword_args(cfg, eng):
+    return {"prefixes": ["TEMPORARY"]}
+
+
+@set_default_schema_on_connection.for_db("postgresql")
+def _postgresql_set_default_schema_on_connection(
+    cfg, dbapi_connection, schema_name
+):
+    existing_autocommit = dbapi_connection.autocommit
+    dbapi_connection.autocommit = True
+    cursor = dbapi_connection.cursor()
+    cursor.execute("SET SESSION search_path='%s'" % schema_name)
+    cursor.close()
+    dbapi_connection.autocommit = existing_autocommit
+
+
+@drop_all_schema_objects_pre_tables.for_db("postgresql")
+def drop_all_schema_objects_pre_tables(cfg, eng):
+    with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
+        for xid in conn.exec_driver_sql(
+            "select gid from pg_prepared_xacts"
+        ).scalars():
+            conn.exec_driver_sql("ROLLBACK PREPARED '%s'" % xid)
+
+
+@drop_all_schema_objects_post_tables.for_db("postgresql")
+def drop_all_schema_objects_post_tables(cfg, eng):
+    from sqlalchemy.dialects import postgresql
+
+    inspector = inspect(eng)
+    with eng.begin() as conn:
+        for enum in inspector.get_enums("*"):
+            conn.execute(
+                postgresql.DropEnumType(
+                    postgresql.ENUM(name=enum["name"], schema=enum["schema"])
+                )
+            )
+
+
+@prepare_for_drop_tables.for_db("postgresql")
+def prepare_for_drop_tables(config, connection):
+    """Ensure there are no locks on the current username/database."""
+
+    result = connection.exec_driver_sql(
+        "select pid, state, wait_event_type, query "
+        # "select pg_terminate_backend(pid), state, wait_event_type "
+        "from pg_stat_activity where "
+        "usename=current_user "
+        "and datname=current_database() and state='idle in transaction' "
+        "and pid != pg_backend_pid()"
+    )
+    rows = result.all()  # noqa
+    if rows:
+        warn_test_suite(
+            "PostgreSQL may not be able to DROP tables due to "
+            "idle in transaction: %s"
+            % ("; ".join(row._mapping["query"] for row in rows))
+        )
+
+
+@upsert.for_db("postgresql")
+def _upsert(
+    cfg, table, returning, *, set_lambda=None, sort_by_parameter_order=False
+):
+    from sqlalchemy.dialects.postgresql import insert
+
+    stmt = insert(table)
+
+    table_pk = inspect(table).selectable
+
+    if set_lambda:
+        stmt = stmt.on_conflict_do_update(
+            index_elements=table_pk.primary_key, set_=set_lambda(stmt.excluded)
+        )
+    else:
+        stmt = stmt.on_conflict_do_nothing()
+
+    stmt = stmt.returning(
+        *returning, sort_by_parameter_order=sort_by_parameter_order
+    )
+    return stmt
+
+
+_extensions = [
+    ("citext", (13,)),
+    ("hstore", (13,)),
+]
+
+
+@post_configure_engine.for_db("postgresql")
+def _create_citext_extension(url, engine, follower_ident):
+    with engine.connect() as conn:
+        for extension, min_version in _extensions:
+            if conn.dialect.server_version_info >= min_version:
+                conn.execute(
+                    text(f"CREATE EXTENSION IF NOT EXISTS {extension}")
+                )
+                conn.commit()