Add PostgresSchema to manage Postgres schemas with

Author: Photonios

docs/source/api_reference.rst

@@ -34,12 +34,17 @@ API Reference
 .. automodule:: psqlextra.indexes
 
    .. autoclass:: UniqueIndex
+
    .. autoclass:: ConditionalUniqueIndex
+
    .. autoclass:: CaseInsensitiveUniqueIndex
 
 .. automodule:: psqlextra.locking
    :members:
 
+.. automodule:: psqlextra.schema
+   :members:
+
 .. automodule:: psqlextra.partitioning
    :members:
 

docs/source/index.rst

@@ -40,6 +40,11 @@ Explore the documentation to learn about all features:
    Support for explicit table-level locks.
 
 
+* :ref:`Creating/dropping schemas <schemas_page>`
+
+    Support for managing Postgres schemas.
+
+
 .. toctree::
    :maxdepth: 2
    :caption: Overview
@@ -54,6 +59,7 @@ Explore the documentation to learn about all features:
    expressions
    annotations
    locking
+   schemas
    settings
    api_reference
    major_releases

docs/source/schemas.rst

@@ -0,0 +1,169 @@
+.. include:: ./snippets/postgres_doc_links.rst
+
+.. _schemas_page:
+
+Schema
+======
+
+The :meth:`~psqlextra.schema.PostgresSchema` class provides basic schema management functionality.
+
+Django does **NOT** support custom schemas. This module does not attempt to solve that problem.
+
+This module merely allows you to create/drop schemas and allow you to execute raw SQL in a schema. It is not attempt at bringing multi-schema support to Django.
+
+
+Reference an existing schema
+----------------------------
+
+.. code-block:: python
+
+   for psqlextra.schema import PostgresSchema
+
+   schema = PostgresSchema("myschema")
+
+   with schema.connection.cursor() as cursor:
+       cursor.execute("SELECT * FROM tablethatexistsinmyschema")
+
+
+Checking if a schema exists
+---------------------------
+
+.. code-block:: python
+
+   for psqlextra.schema import PostgresSchema
+
+   schema = PostgresSchema("myschema")
+   if PostgresSchema.exists("myschema"):
+       print("exists!")
+   else:
+       print('does not exist!")
+
+
+Creating a new schema
+---------------------
+
+With a custom name
+******************
+
+.. code-block:: python
+
+   for psqlextra.schema import PostgresSchema
+
+   # will raise an error if the schema already exists
+   schema = PostgresSchema.create("myschema")
+
+
+Re-create if necessary with a custom name
+*****************************************
+
+.. warning::
+
+   If the schema already exists and it is non-empty or something is referencing it, it will **NOT** be dropped. Specify ``cascade=True`` to drop all of the schema's contents and **anything referencing it**.
+
+.. code-block:: python
+
+   for psqlextra.schema import PostgresSchema
+
+   # will drop existing schema named `myschema` if it
+   # exists and re-create it
+   schema = PostgresSchema.drop_and_create("myschema")
+
+   # will drop the schema and cascade it to its contents
+   # and anything referencing the schema
+   schema = PostgresSchema.drop_and_create("otherschema", cascade=True)
+
+
+With a random name
+******************
+
+.. code-block:: python
+
+   for psqlextra.schema import PostgresSchema
+
+   # schema name will be "myprefix_<timestamp>"
+   schema = PostgresSchema.create_random("myprefix")
+   print(schema.name)
+
+
+Temporary schema with random name
+*********************************
+
+Use the :meth:`~psqlextra.schema.postgres_temporary_schema` context manager to create a schema with a random name. The schema will only exist within the context manager.
+
+By default, the schema is not dropped if an exception occurs in the context manager. This prevents unexpected data loss. Specify ``drop_on_throw=True`` to drop the schema if an exception occurs.
+
+Without an outer transaction, the temporary schema might not be dropped when your program is exits unexpectedly (for example; if it is killed with SIGKILL). Wrap the creation of the schema in a transaction to make sure the schema is cleaned up when an error occurs or your program exits suddenly.
+
+.. warning::
+
+   By default, the drop will fail if the schema is not empty or there is anything referencing the schema.  Specify ``cascade=True`` to drop all of the schema's contents and **anything referencing it**.
+
+.. note::
+
+
+.. code-block:: python
+
+   for psqlextra.schema import postgres_temporary_schema
+
+   with postgres_temporary_schema("myprefix") as schema:
+       pass
+
+   with postgres_temporary_schema("otherprefix", drop_on_throw=True) as schema:
+       raise ValueError("drop it like it's hot")
+
+   with postgres_temporary_schema("greatprefix", cascade=True) as schema:
+       with schema.connection.cursor() as cursor:
+           cursor.execute(f"CREATE TABLE {schema.name} AS SELECT 'hello'")
+
+   with postgres_temporary_schema("amazingprefix", drop_on_throw=True, cascade=True) as schema:
+       with schema.connection.cursor() as cursor:
+           cursor.execute(f"CREATE TABLE {schema.name} AS SELECT 'hello'")
+
+       raise ValueError("oops")
+
+Deleting a schema
+-----------------
+
+Any schema can be dropped, including ones not created by :class:`~psqlextra.schema.PostgresSchema`.
+
+The ``public`` schema cannot be dropped. This is a Postgres built-in and it is almost always a mistake to drop it. A :class:`~django.core.exceptions.SuspiciousOperation` erorr will be raised if you attempt to drop the ``public`` schema.
+
+.. warning::
+
+   By default, the drop will fail if the schema is not empty or there is anything referencing the schema.  Specify ``cascade=True`` to drop all of the schema's contents and **anything referencing it**.
+
+.. code-block:: python
+
+   for psqlextra.schema import PostgresSchema
+
+   schema = PostgresSchema.drop("myprefix")
+   schema = PostgresSchema.drop("myprefix", cascade=True)
+
+
+Executing queries within a schema
+---------------------------------
+
+By default, a connection operates in the ``public`` schema. The schema offers a connection scoped to that schema that sets the Postgres ``search_path`` to only search within that schema.
+
+.. warning::
+
+   This can be abused to manage Django models in a custom schema. This is not a supported workflow and there might be unexpected issues from attempting to do so.
+
+.. warning::
+
+   Do not pass the connection to a different thread. It is **NOT** thread safe.
+
+.. code-block:: python
+
+   from psqlextra.schema import PostgresSchema
+
+   schema = PostgresSchema.create("myschema")
+
+   with schema.connection.cursor() as cursor:
+       # table gets created within the `myschema` schema, without
+       # explicitly specifying the schema name
+       cursor.execute("CREATE TABLE mytable AS SELECT 'hello'")
+
+    with schema.connection.schema_editor() as schema_editor:
+        # creates a table for the model within the schema
+        schema_editor.create_model(MyModel)

psqlextra/backend/introspection.py

@@ -69,7 +69,8 @@ def get_partitioned_tables(
     ) -> PostgresIntrospectedPartitonedTable:
         """Gets a list of partitioned tables."""
 
-        sql = """
+        cursor.execute(
+            """
             SELECT
                 pg_class.relname,
                 pg_partitioned_table.partstrat
@@ -80,8 +81,7 @@ def get_partitioned_tables(
             ON
                 pg_class.oid = pg_partitioned_table.partrelid
         """
-
-        cursor.execute(sql)
+        )
 
         return [
             PostgresIntrospectedPartitonedTable(
@@ -191,6 +191,19 @@ def get_partition_key(self, cursor, table_name: str) -> List[str]:
     def get_columns(self, cursor, table_name: str):
         return self.get_table_description(cursor, table_name)
 
+    def get_schema_list(self, cursor) -> List[str]:
+        """A flat list of available schemas."""
+
+        sql = """
+            SELECT
+                schema_name
+            FROM
+                information_schema.schemata
+        """
+
+        cursor.execute(sql, tuple())
+        return [name for name, in cursor.fetchall()]
+
     def get_constraints(self, cursor, table_name: str):
         """Retrieve any constraints or keys (unique, pk, fk, check, index)
         across one or more columns.

psqlextra/backend/schema.py

@@ -45,6 +45,9 @@ class PostgresSchemaEditor(SchemaEditor):
     sql_reset_table_storage_setting = "ALTER TABLE %s RESET (%s)"
 
     sql_alter_table_schema = "ALTER TABLE %s SET SCHEMA %s"
+    sql_create_schema = "CREATE SCHEMA %s"
+    sql_delete_schema = "DROP SCHEMA %s"
+    sql_delete_schema_cascade = "DROP SCHEMA %s CASCADE"
 
     sql_create_view = "CREATE VIEW %s AS (%s)"
     sql_replace_view = "CREATE OR REPLACE VIEW %s AS (%s)"
@@ -84,6 +87,21 @@ def __init__(self, connection, collect_sql=False, atomic=True):
         self.deferred_sql = []
         self.introspection = PostgresIntrospection(self.connection)
 
+    def create_schema(self, name: str) -> None:
+        """Creates a Postgres schema."""
+
+        self.execute(self.sql_create_schema % self.quote_name(name))
+
+    def delete_schema(self, name: str, cascade: bool) -> None:
+        """Drops a Postgres schema."""
+
+        sql = (
+            self.sql_delete_schema
+            if not cascade
+            else self.sql_delete_schema_cascade
+        )
+        self.execute(sql % self.quote_name(name))
+
     def create_model(self, model: Type[Model]) -> None:
         """Creates a new model."""
 

psqlextra/error.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+import psycopg2
+
+from django import db
+
+
+def extract_postgres_error(error: db.Error) -> Optional[psycopg2.Error]:
+    """Extracts the underlying :see:psycopg2.Error from the specified Django
+    database error.
+
+    As per PEP-249, Django wraps all database errors in its own
+    exception. We can extract the underlying database error by examaning
+    the cause of the error.
+    """
+
+    if not isinstance(error.__cause__, psycopg2.Error):
+        return None
+
+    return error.__cause__