python 简单的帮助器，支持Postgresql 10的声明性分区功能。

Posted 2021-05-09

import textwrap
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.ddl import CreateTable


@compiles(CreateTable, "postgresql")
def pg10_partition_compiler(element, compiler, **kw):
    """
    Simple helper that enables Postgresql 10's declarative partitioning features.
        partition_by : str = Plain sql string that defines partition rules for the parent table.
            eg. partition_by='RANGE (logdate)' or partition_by='LIST (category)'
        partition_of : table = SqlAlchemy table object for the parent table.
            eg. partition_of=MyTable or partition_of=yourmodel.__table__
        for_values : str = Plain sql string for the child table to partition.
            eg. for_values='IN (others)' or for_values='FROM ('2016-07-01') TO ('2016-08-01')'
    """
    compiled_sql = compiler.visit_create_table(element, **kw)
    if set(element.element.kwargs.keys()).issuperset({"partition_by", "for_values", "partition_of"}):
        raise ValueError("You can't define partition_by and partition_of & for_values at the same table.")
    elif "partition_by" in element.element.kwargs:
        compiled_sql += "PARTITION BY " + textwrap.dedent(element.element.kwargs["partition_by"]).strip()
    elif "partition_of" in element.element.kwargs and "for_values" in element.element.kwargs:
        compiled_sql = " ".join([
            compiled_sql[:compiled_sql.find("(")],
            "PARTITION OF " + textwrap.dedent(element.element.kwargs["partition_of"].name).strip(),
            compiled_sql[compiled_sql.find(")") + 1:],  # omit (\n) part from the query
            "FOR VALUES " + textwrap.dedent(element.element.kwargs["for_values"]).strip()
        ])
    return compiled_sql