~launchpad-pqm/launchpad/devel

#!/usr/bin/python

#

# Copyright 2009 Canonical Ltd.  This software is licensed under the

# GNU Affero General Public License version 3 (see the file LICENSE).

"""

dropdb only more so.

Cut off access, slaughter connections and burn the database to the ground.

"""

# Nothing but system installed libraries - this script sometimes

# gets installed standalone with no Launchpad tree available.

from distutils.version import LooseVersion

import sys

import time

import psycopg2

import psycopg2.extensions

from signal import SIGTERM, SIGQUIT, SIGKILL, SIGINT

from optparse import OptionParser

def connect(dbname='template1'):

    """Connect to the database, returning the DB-API connection."""

    if options.user is not None:

        return psycopg2.connect("dbname=%s user=%s" % (dbname, options.user))

    else:

        return psycopg2.connect("dbname=%s" % dbname)

def send_signal(database, signal):

    con = connect()

    con.set_isolation_level(1) # READ COMMITTED. We rollback changes we make.

    cur = con.cursor()

    # Install PL/PythonU if it isn't already.

    cur.execute("SELECT TRUE FROM pg_language WHERE lanname = 'plpythonu'")

    if cur.fetchone() is None:

        cur.execute('CREATE LANGUAGE "plpythonu"')

    # Create a stored procedure to kill a backend process.

    qdatabase = str(psycopg2.extensions.QuotedString(database))

    cur.execute("""

        CREATE OR REPLACE FUNCTION _pgmassacre_killall(integer)

        RETURNS Boolean AS $$

        import os

        signal = args[0]

        for row in plpy.execute('''

            SELECT procpid FROM pg_stat_activity WHERE datname=%(qdatabase)s

                AND procpid != pg_backend_pid()

            '''):

            try:

                os.kill(row['procpid'], signal)

            except OSError:

                pass

        else:

            return False

        return True

        $$ LANGUAGE plpythonu

        """ % vars())

    cur.execute("SELECT _pgmassacre_killall(%(signal)s)", vars())

    con.rollback()

    con.close()

def rollback_prepared_transactions(database):

    """Rollback any prepared transactions.

    PostgreSQL will refuse to drop a database with outstanding prepared

    transactions.

    """

    con = connect(database)

    con.set_isolation_level(0) # Autocommit so we can ROLLBACK PREPARED.

    cur = con.cursor()

    # Get a list of outstanding prepared transactions.

    cur.execute(

            "SELECT gid FROM pg_prepared_xacts WHERE database=%(database)s",

            vars())

    xids = [row[0] for row in cur.fetchall()]

    for xid in xids:

        cur.execute("ROLLBACK PREPARED %(xid)s", vars())

    con.close()

def still_open(database, max_wait=10):

    """Return True if there are still open connections, apart from our own.

    Waits a while to ensure that connections shutting down have a chance to.

    """

    con = connect()

    con.set_isolation_level(0) # Autocommit.

    cur = con.cursor()

    # Wait for up to 10 seconds, returning True if all backends are gone.

    start = time.time()

    while time.time() < start + max_wait:

        cur.execute("""

            SELECT TRUE FROM pg_stat_activity

            WHERE

                datname=%(database)s

                AND procpid != pg_backend_pid()

            LIMIT 1

            """, vars())

        if cur.fetchone() is None:

            return False

        time.sleep(0.6) # Stats only updated every 500ms.

    con.close()

    return True

def massacre(database):

    con = connect()

    con.set_isolation_level(0) # Autocommit

    cur = con.cursor()

    # Allow connections to the doomed database if something turned this off,

    # such as an aborted run of this script.

    cur.execute(

        "UPDATE pg_database SET datallowconn=TRUE WHERE datname=%s",

        [database])

    # Rollback prepared transactions.

    rollback_prepared_transactions(database)

    try:

        # Stop connections to the doomed database.

        cur.execute(

            "UPDATE pg_database SET datallowconn=FALSE WHERE datname=%s",

            [database])

        con.close()

        # Terminate current statements.

        send_signal(database, SIGINT)

        # Shutdown current connections normally.

        if still_open(database, 1):

            send_signal(database, SIGTERM)

        # Shutdown current connections immediately.

        if still_open(database):

            send_signal(database, SIGQUIT)

        # Shutdown current connections nastily.

        if still_open(database):

            send_signal(database, SIGKILL)

        if still_open(database):

            print >> sys.stderr, (

                    "Unable to kill all backends! Database not destroyed.")

            return 9

        # Destroy the database.

        con = connect()

        # AUTOCOMMIT required to execute commands like DROP DATABASE.

        con.set_isolation_level(0)

        cur = con.cursor()

        cur.execute("DROP DATABASE %s" % database) # Not quoted.

        con.close()

        return 0

    finally:

        # In case something messed up, allow connections again so we can

        # inspect the damage.

        con = connect()

        con.set_isolation_level(0)

        cur = con.cursor()

        cur.execute(

                "UPDATE pg_database SET datallowconn=TRUE WHERE datname=%s",

                [database])

        con.close()

def rebuild(database, template):

    if still_open(template, 20):

        print >> sys.stderr, (

            "Giving up waiting for connections to %s to drop." % template)

        report_open_connections(template)

        return 10

    start = time.time()

    now = start

    error_msg = None

    con = connect()

    con.set_isolation_level(0) # Autocommit required for CREATE DATABASE.

    create_db_cmd = """

        CREATE DATABASE %s WITH ENCODING='UTF8' TEMPLATE=%s

        """ % (database, template)

    # 8.4 allows us to create empty databases with a different locale

    # to template1 by using the template0 database as a template.

    # We make use of this feature so we don't have to care what locale

    # was used to create the database cluster rather than requiring it

    # to be rebuilt in the C locale.

    if pg_version >= LooseVersion("8.4.0") and template == "template0":

        create_db_cmd += "LC_COLLATE='C' LC_CTYPE='C'"

    while now < start + 20:

        cur = con.cursor()

        try:

            cur.execute(create_db_cmd)

            con.close()

            return 0

        except psycopg2.Error, exception:

            error_msg = str(exception)

        time.sleep(0.6) # Stats only updated every 500ms.

        now = time.time()

    con.close()

    print >> sys.stderr, "Unable to recreate database: %s" % error_msg

    return 11

def report_open_connections(database):

    con = connect()

    cur = con.cursor()

    cur.execute("""

        SELECT usename, datname, count(*)

        FROM pg_stat_activity

        WHERE procpid != pg_backend_pid()

        GROUP BY usename, datname

        ORDER BY datname, usename

        """, [database])

    for usename, datname, num_connections in cur.fetchall():

        print >> sys.stderr, "%d connections by %s to %s" % (

            num_connections, usename, datname)

    con.close()

options = None

pg_version = None # LooseVersion - Initialized in main()

def main():

    parser = OptionParser("Usage: %prog [options] DBNAME")

    parser.add_option("-U", "--user", dest="user", default=None,

        help="Connect as USER", metavar="USER")

    parser.add_option("-t", "--template", dest="template", default=None,

        help="Recreate database using DBNAME as a template database.",

        metavar="DBNAME")

    global options

    (options, args) = parser.parse_args()

    if len(args) != 1:

        parser.error('Must specify one, and only one, database to destroy')

    database = args[0]

    # Don't be stupid protection.

    if database in ('template1', 'template0'):

        parser.error(

            "Running this script against template1 or template0 is nuts.")

    con = connect()

    cur = con.cursor()

    # Store the database version for version specific code.

    global pg_version

    cur.execute("show server_version")

    pg_version = LooseVersion(cur.fetchone()[0])

    # Ensure the template database exists.

    if options.template is not None:

        cur.execute(

            "SELECT TRUE FROM pg_database WHERE datname=%s",

            [options.template])

        if cur.fetchone() is None:

            parser.error(

                "Template database %s does not exist." % options.template)

    # If the database doesn't exist, no point attempting to drop it.

    cur.execute("SELECT TRUE FROM pg_database WHERE datname=%s", [database])

    db_exists = cur.fetchone() is not None

    con.close()

    if db_exists:

        rv = massacre(database)

        if rv != 0:

            print >> sys.stderr, "Fail %d" % rv

            return rv

    if options.template is not None:

        return rebuild(database, options.template)

    else:

        return 0

if __name__ == '__main__':

    sys.exit(main())