| #!/bin/env python |
| # coding: UTF-8 |
| |
| # PsqlWrapper |
| # |
| # Copyright(c) 2015 Uptime Technologies, LLC. |
| |
| import log |
| import os |
| import re |
| import subprocess |
| import sys |
| |
| def get_host(host=None): |
| if host is None: |
| if os.environ.get("PGHOST") is not None: |
| h = os.environ.get("PGHOST") |
| else: |
| h = 'localhost' |
| else: |
| h = host |
| |
| return h |
| |
| def get_port(port=None): |
| if port is None: |
| if os.environ.get("PGPORT") is not None: |
| p = int(os.environ.get("PGPORT")) |
| else: |
| p = 5432 |
| else: |
| p = int(port) |
| |
| return p |
| |
| def get_username(username=None): |
| if username is None: |
| if os.environ.get("PGUSER") is not None: |
| u = os.environ.get("PGUSER") |
| else: |
| u = os.environ.get("USER") |
| else: |
| u = username |
| |
| return u |
| |
| def get_dbname(dbname=None, username=None): |
| if dbname is None: |
| if os.environ.get("PGDATABASE") is not None: |
| d = os.environ.get("PGDATABASE") |
| else: |
| d = username |
| else: |
| d = dbname |
| |
| return d |
| |
| class PsqlWrapper: |
| version = None |
| |
| def __init__(self, host, port, username, dbname, on_error_stop=False, debug=False): |
| if debug is True: |
| log.setLevel(log.DEBUG) |
| |
| self.host = get_host(host) |
| self.port = get_port(port) |
| self.username = get_username(username) |
| self.dbname = get_dbname(dbname, self.username) |
| |
| log.debug("host: " + self.host) |
| log.debug("port: " + str(self.port)) |
| log.debug("user: " + self.username) |
| log.debug("dbname: " + self.dbname) |
| |
| self.on_error_stop = on_error_stop |
| |
| def get_version(self): |
| if self.version is None: |
| rs = self.execute_query('select version()') |
| # print(rs[1]) |
| m = re.search('PostgreSQL (\d+\.\d+)', rs[1][0]) |
| |
| self.version = float(m.group(1)) |
| |
| return self.version |
| |
| def execute_query(self, query, ignore_error=False): |
| cmd = "psql -A" |
| cmd = cmd + " -h " + self.host |
| cmd = cmd + " -p " + str(self.port) |
| cmd = cmd + " -U " + self.username |
| cmd = cmd + " -d " + self.dbname |
| if self.on_error_stop is True: |
| cmd = cmd + " --set=ON_ERROR_STOP=on " |
| |
| log.debug(cmd) |
| |
| p = subprocess.Popen(cmd, shell=True, |
| stdin=subprocess.PIPE, stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE) |
| stdout_data, stderr_data = p.communicate(input=query) |
| p.returncode, stdout_data, stderr_data |
| |
| self.stderr_data = stderr_data |
| |
| if p.returncode != 0: |
| if ignore_error is True: |
| log.debug("return code: " + str(p.returncode)) |
| log.debug(stderr_data) |
| return None |
| else: |
| log.debug("return code: " + str(p.returncode)) |
| log.error("Failed to execute psql. (" + cmd + ")") |
| if stderr_data is not None: |
| log.error(stderr_data.replace('\n','')) |
| sys.exit(1) |
| |
| rs = [] |
| lines = re.split('\n', stdout_data)[:-1] |
| for line in lines: |
| rs.append(re.split('\|', line)) |
| |
| return rs |
| |
| def print_result(self, rs): |
| cols = len(rs[0]) |
| size = [] |
| |
| header = True |
| for r in rs: |
| if len(r) != cols: |
| continue |
| |
| i = 0 |
| for c in r: |
| if header is True: |
| size.append(len(c)) |
| elif len(c) > size[i]: |
| size[i] = len(c) |
| |
| # print '[' + str(i) + ']' + c |
| i = i + 1 |
| header = False |
| |
| # print size |
| |
| header = True |
| sep2 = None |
| for r in rs: |
| if len(r) != cols: |
| continue |
| |
| i = 0 |
| out = '|' |
| sep = '+' |
| for c in r: |
| if header is True: |
| out = out + ' ' + str(c).center(size[i]) + ' |' |
| sep = sep + '-' + '-' * size[i] + '-+' |
| elif re.match('^\d+$', str(c)): |
| out = out + ' ' + str(c).rjust(size[i]) + ' |' |
| else: |
| out = out + ' ' + str(c).ljust(size[i]) + ' |' |
| i = i + 1 |
| |
| if header is True: |
| sep2 = sep |
| print sep2 |
| |
| print out |
| |
| if header is True: |
| print sep2 |
| header = False |
| |
| print sep2 |