Add PostgresSchema to manage Postgres schemas with

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__