bin/pt-set-tablespace - postgres-toolkit - Git at Google

 #!/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)
	#!/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)