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