Nsryjdtyk

I'm importing millions of rows into a PostgreSQL database using a foreign data wrapper. To speed this up I've written this context manager to DROP all the indexes and constraints (except the primary key) on the target table for the duration of the data move and reinstate them afterwards:

import logging

import types

import typing

from psycopg2 import _psycopg, sql





class RawTable():

    def __init__(self, cursor: _psycopg.cursor, table_name: str, schema_name: str = 'public'):

        self.cursor = cursor

        self.schema_name = schema_name

        self.table_name = table_name



        self.constraints: typing.List[typing.Tuple[str, str]] = 

        self.index_definitions: typing.List[str] = 



    def __enter__(self) -> None:

        self.store_and_delete_constraints()

        self.store_and_delete_indexes()



    def store_and_delete_constraints(self) -> None:

        self.cursor.execute(

            """

            SELECT pg_constraint.conname,

                   pg_get_constraintdef(pg_constraint.oid)

              FROM pg_constraint

              JOIN pg_class ON pg_constraint.conrelid = pg_class.oid

              JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace

             WHERE pg_constraint.contype <> 'p' AND

                   pg_class.relname = %s AND

                   pg_namespace.nspname = %s

            """,

            (self.table_name, self.schema_name)

        )

        for (name, constraint_sql) in self.cursor.fetchall():

            self.constraints.append((name, constraint_sql))

            logging.debug("Dropping constraint {}".format(name))

            self.cursor.execute(

                sql.SQL('ALTER TABLE {}.{} DROP CONSTRAINT {}').format(

                    sql.Identifier(self.schema_name),

                    sql.Identifier(self.table_name),

                    sql.Identifier(name)

                )

            )



    def store_and_delete_indexes(self) -> None:

        self.cursor.execute(

            """

            SELECT index_class.relname,

                   pg_get_indexdef(pg_index.indexrelid)

              FROM pg_index

              JOIN pg_class AS index_class ON pg_index.indexrelid = index_class.oid

              JOIN pg_class AS table_class ON pg_index.indrelid = table_class.oid

              JOIN pg_namespace ON pg_namespace.oid = index_class.relnamespace

             WHERE NOT pg_index.indisprimary AND

                   pg_index.indisvalid AND

                   table_class.relname = %s AND

                   pg_namespace.nspname = %s

            """,

            (self.table_name, self.schema_name)

        )

        for (name, index_sql) in self.cursor.fetchall():

            self.index_definitions.append(index_sql)

            logging.debug("Dropping index {}".format(name))

            self.cursor.execute(sql.SQL('DROP INDEX {}').format(sql.Identifier(name)))



    def __exit__(

            self,

            exception_type: typing.Type[BaseException],

            exception: BaseException,

            traceback: types.TracebackType

    ) -> bool:

        self.restore_constraints()

        self.restore_indexes()

        return exception_type is None



    def restore_constraints(self) -> None:

        logging.debug("Restoring {} constraint(s)".format(len(self.constraints)))

        constraint_definitions = [

            sql.SQL("ADD CONSTRAINT {} {}").format(

                sql.Identifier(name),

                sql.SQL(constraint_sql)

            ) for name, constraint_sql in self.constraints

        ]

        self.cursor.execute(

            sql.SQL("ALTER TABLE {}.{} {}").format(

                sql.Identifier(self.schema_name),

                sql.Identifier(self.table_name),

                sql.SQL(', ').join(constraint_definitions)

            )

        )



    def restore_indexes(self) -> None:

        logging.debug("Restoring {} index(es)".format(len(self.index_definitions)))

        for index_definition in self.index_definitions:

            self.cursor.execute(index_definition)

Example use:

with connection.cursor() as cursor, RawTable(cursor, 'target_table'):

    cursor.execute('INSERT INTO target_table SELECT … FROM huge_table')

(Log level is expected to be set up elsewhere, hence the lack of a logging.basicConfig or equivalent call.)