bin/pt-set-tablespace

#!/usr/bin/env python
# coding: UTF-8

# pt-set-tablespace
#
# Copyright(c) 2015 Uptime Technologies, LLC.

import sys, os
libpath = os.path.abspath(os.path.dirname(sys.argv[0]) + "/../lib")
sys.path.append(libpath)

import re
import getopt
import subprocess
import log
import PsqlWrapper

class SetTablespace:
    alter_stmt = []

    def build_where_clause(self, where_clause, cond):
        if where_clause is None:
            where_clause = " where "
        else:
            where_clause = where_clause + " and "
        where_clause = where_clause + " " + cond
        return where_clause

    def __init__(self, psql, owner, schema, table, tablespace, debug=False):
        self.debug = debug

        self.psql   = psql
        self.owner  = owner
        self.schema = schema
        self.table  = table

        self.tablespace = tablespace

    def build_where(self):
        where_clause = None

        if self.schema is not None:
            where_clause = self.build_where_clause(where_clause, "t.schemaname = '" + self.schema + "'")

        if self.owner is not None:
            where_clause = self.build_where_clause(where_clause, "u.usename = '" + self.owner + "'")

        if self.table is not None:
            where_clause = self.build_where_clause(where_clause, "t.relname = '" + self.table + "'")

        if where_clause is None:
            where_clause = ''

        return where_clause

    def alter_table(self):
        where_clause = self.build_where()

        self.query = ' \
select \
       t.relid, \
       u.usename, \
       t.schemaname, \
       t.relname, \
       s.spcname \
  from \
       pg_stat_user_tables t left outer join pg_class c \
                on t.relid = c.oid \
                             left outer join pg_tablespace s \
                on c.reltablespace = s.oid \
                             left outer join pg_user u \
                on c.relowner = u.usesysid \
%s \
 order by \
    2,3,4 \
;' % (where_clause)

        log.debug(self.query)

        rs = self.psql.execute_query(self.query)
        
        row = 0
        for r in rs:
            if row > 0 and len(r) == 5:
                self.alter_stmt.append("ALTER TABLE \"%s\".\"%s\" SET TABLESPACE \"%s\";" % (r[2], r[3], self.tablespace))
            row = row + 1

        if self.debug is True:
            self.psql.print_result(rs)

    def alter_index(self):
        where_clause = self.build_where()

        self.query = ' \
select \
       t.indexrelid, \
       u.usename, \
       t.schemaname, \
       t.relname, \
       t.indexrelname, \
       s.spcname \
  from \
       pg_stat_user_indexes t left outer join pg_class c \
                on t.indexrelid = c.oid \
                             left outer join pg_tablespace s \
                on c.reltablespace = s.oid \
                             left outer join pg_user u \
                on c.relowner = u.usesysid \
%s \
;' % (where_clause)

        log.debug(self.query)

        rs = self.psql.execute_query(self.query)
        
        row = 0
        for r in rs:
            if row > 0 and len(r) == 6:
                self.alter_stmt.append("ALTER INDEX \"%s\".\"%s\" SET TABLESPACE \"%s\";" % (r[2], r[4], self.tablespace))
            row = row + 1

        if self.debug is True:
            self.psql.print_result(rs)

    def apply(self, do_apply):
        self.alter_table()
        self.alter_index()

        if do_apply is True:
            log.info("Applying ALTER TABLE/INDEX...")
        else:
            log.info("Dry-run mode:")

        failed    = 0
        succeeded = 0

        for s in self.alter_stmt:
            if do_apply is True:
                log.debug(s)

                rs = self.psql.execute_query(s)

                log.debug(str(rs))

                if rs is None or len(rs) == 0:
                    log.error(str(self.psql.stderr_data).replace("ERROR: ", "").replace("\n", ""))
                    log.error("  " + s)
                    failed = failed + 1

                elif rs[0][0] != 'ALTER TABLE' and rs[0][0] != 'ALTER INDEX':
                    log.error(str(self.psql.stderr_data).replace("ERROR: ", "").replace("\n", ""))
                    log.error("  " + s)
                    failed = failed + 1

                else:
                    succeeded = succeeded + 1
                    log.debug("Succeeded")

            else:
                log.info(s)

        if do_apply is True:
            log.info("%d tables/indexes moved. %d failed." % (succeeded, failed))

        if failed > 0:
            return False

        return True


def tablespace_size(spcloc):
    log.debug("location: " + spcloc)

    if spcloc is None or len(spcloc) == 0:
        return ['', '']

    cmd = "df -khPl " + re.sub('/[^/]+$', '', spcloc) + " | grep ^/ | awk '{ print $5 \" \" $4 }'"

    p = subprocess.Popen([cmd], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE, shell=True)
    p.wait()
    
    if p.returncode != 0:
        log.error("Cannot execute df command.")
        sys.exit(1)

    line = None
    for l in p.stdout.readlines():
        if line is None:
            line = l.replace('\n', '')
        break

    return str(line).split(' ')


def list_tablespace(psql):
    if psql.get_version() <= 9.1:
        query = ' \
select s.oid as "OID", \
       usename as "OWNER", \
       spcname as "TABLESPACE", \
       spclocation as "LOCATION", \
       null as "USE%", \
       null as "AVAIL" \
  from pg_tablespace s left outer join pg_user u \
           on s.spcowner = u.usesysid; \
';
    else:
        query = ' \
select s.oid as "OID", \
       usename as "OWNER", \
       spcname as "TABLESPACE", \
       pg_tablespace_location(s.oid) as "LOCATION", \
       null as "USE%", \
       null as "AVAIL" \
  from pg_tablespace s left outer join pg_user u \
           on s.spcowner = u.usesysid; \
';

    rs = psql.execute_query(query)

    for r in rs:
        if len(r) >= 3 and len(r[3]) > 0 and r[3] != 'LOCATION':
            a = tablespace_size(r[3])
            r[4] = a[0]
            r[5] = a[1]

    psql.print_result(rs)

def usage():
    print ""
    print "Usage: " + os.path.basename(sys.argv[0]) + " [option...] [tablespace]"
    print ""
    print "Options:"
    print "    -h, --host=HOSTNAME        Host name of the postgres server"
    print "    -p, --port=PORT            Port number of the postgres server"
    print "    -U, --username=USERNAME    User name to connect"
    print "    -d, --dbname=DBNAME        Database name to connect"
    print ""
    print "    -o, --owner=STRING         Owner name"
    print "    -n, --schema=STRING        Schema name"
    print "    -t, --table=STRING         Table name"
    print ""
    print "    -l, --list                 List table spaces"
    print "    --apply                    Apply change(s)"
    print ""
    print "    --help                     Print this help."
    print ""

if __name__ == "__main__":
    try:
        opts, args = getopt.getopt(sys.argv[1:], "h:p:U:d:o:n:lt:i:u",
                                   ["help", "debug", "host=", "port=", "username=", "dbname=",
                                    "owner=", "schema=", "table=", "list", "apply"])
    except getopt.GetoptError, err:
        print str(err)
        usage()
        sys.exit(2)

    host     = None
    port     = None
    username = None
    dbname   = None

    owner    = None
    schema   = None
    table    = None

    do_list  = False
    do_apply = False

    debug    = None

    for o, a in opts:
        if o in ("-h", "--host"):
            host = a
        elif o in ("-p", "--port"):
            port = int(a)
        elif o in ("-U", "--username"):
            username = a
        elif o in ("-d", "--dbname"):
            dbname = a
        elif o in ("-o", "--owner"):
            owner = a
        elif o in ("-n", "--schema"):
            schema = a
        elif o in ("-t", "--table"):
            table = a
        elif o in ("-l", "--list"):
            do_list = True
        elif o in ("--apply"):
            do_apply = True
        elif o in ("--debug"):
            debug = True
        elif o in ("--help"):
            usage()
            sys.exit(0)
        else:
            print "unknown option: " + o + "," + a
            sys.exit(1)

    p = PsqlWrapper.PsqlWrapper(host=host, port=port, username=username, dbname=dbname, debug=debug)

    # list tablespaces and exit.
    if do_list is True:
        list_tablespace(p)
        sys.exit(0)

    if len(args) == 1:
        tablespace = args[0]
    else:
        usage()
        sys.exit(0)

    cmd = SetTablespace(p, owner, schema, table, tablespace, debug=debug)

    if cmd.apply(do_apply) is False:
        sys.exit(1)

    sys.exit(0)