2015-01-05 00:20:52 +00:00
|
|
|
import logging
|
2014-11-22 07:43:11 +00:00
|
|
|
import psycopg2
|
2015-01-09 09:56:22 +00:00
|
|
|
import psycopg2.extras
|
2015-01-06 21:52:43 +00:00
|
|
|
import psycopg2.extensions
|
2015-01-08 10:56:28 +00:00
|
|
|
from collections import defaultdict
|
2014-12-08 08:43:21 +00:00
|
|
|
from .packages import pgspecial
|
2014-11-22 07:43:11 +00:00
|
|
|
|
2015-01-05 00:20:52 +00:00
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
2015-01-06 21:52:43 +00:00
|
|
|
# Cast all database input to unicode automatically.
|
|
|
|
# See http://initd.org/psycopg/docs/usage.html#unicode-handling for more info.
|
|
|
|
psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
|
|
|
|
psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
|
|
|
|
|
2015-01-09 09:56:22 +00:00
|
|
|
# When running a query, make pressing CTRL+C raise a KeyboardInterrupt
|
|
|
|
psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
|
|
|
|
|
2014-12-13 02:28:55 +00:00
|
|
|
def _parse_dsn(dsn, default_user, default_password, default_host,
|
|
|
|
default_port):
|
|
|
|
"""
|
2014-12-14 22:18:26 +00:00
|
|
|
This function parses a postgres url to get the different components.
|
|
|
|
|
2014-12-13 02:28:55 +00:00
|
|
|
"""
|
|
|
|
|
|
|
|
user = password = host = port = dbname = None
|
|
|
|
|
2014-12-14 22:18:26 +00:00
|
|
|
if dsn.startswith('postgres://'): # Check if the string is a database url.
|
|
|
|
dsn = dsn[len('postgres://'):]
|
|
|
|
elif dsn.startswith('postgresql://'):
|
|
|
|
dsn = dsn[len('postgresql://'):]
|
2014-12-13 02:28:55 +00:00
|
|
|
|
2014-12-14 22:18:26 +00:00
|
|
|
if '/' in dsn:
|
2014-12-13 02:28:55 +00:00
|
|
|
host, dbname = dsn.split('/', 1)
|
|
|
|
if '@' in host:
|
|
|
|
user, _, host = host.partition('@')
|
|
|
|
if ':' in host:
|
|
|
|
host, _, port = host.partition(':')
|
2014-12-14 22:18:26 +00:00
|
|
|
if user and ':' in user:
|
2014-12-13 02:28:55 +00:00
|
|
|
user, _, password = user.partition(':')
|
|
|
|
|
|
|
|
user = user or default_user
|
|
|
|
password = password or default_password
|
|
|
|
host = host or default_host
|
|
|
|
port = port or default_port
|
|
|
|
dbname = dbname or dsn
|
|
|
|
|
2015-01-05 00:20:52 +00:00
|
|
|
_logger.debug('Parsed connection params:'
|
|
|
|
'dbname: %r, user: %r, password: %r, host: %r, port: %r',
|
|
|
|
dbname, user, password, host, port)
|
|
|
|
|
2014-12-13 02:28:55 +00:00
|
|
|
return (dbname, user, password, host, port)
|
|
|
|
|
2014-11-22 07:43:11 +00:00
|
|
|
class PGExecute(object):
|
2014-11-23 23:02:05 +00:00
|
|
|
|
2014-12-08 08:43:21 +00:00
|
|
|
tables_query = '''SELECT c.relname as "Name" FROM pg_catalog.pg_class c
|
|
|
|
LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace WHERE
|
|
|
|
c.relkind IN ('r','') AND n.nspname <> 'pg_catalog' AND n.nspname <>
|
|
|
|
'information_schema' AND n.nspname !~ '^pg_toast' AND
|
|
|
|
pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1;'''
|
2014-11-23 23:02:05 +00:00
|
|
|
|
2015-01-08 10:56:28 +00:00
|
|
|
columns_query = '''SELECT table_name, column_name FROM information_schema.columns'''
|
2014-11-23 23:02:24 +00:00
|
|
|
|
2015-01-05 00:04:21 +00:00
|
|
|
databases_query = """SELECT d.datname as "Name",
|
|
|
|
pg_catalog.pg_get_userbyid(d.datdba) as "Owner",
|
|
|
|
pg_catalog.pg_encoding_to_char(d.encoding) as "Encoding",
|
|
|
|
d.datcollate as "Collate",
|
|
|
|
d.datctype as "Ctype",
|
|
|
|
pg_catalog.array_to_string(d.datacl, E'\n') AS "Access privileges"
|
|
|
|
FROM pg_catalog.pg_database d
|
|
|
|
ORDER BY 1;"""
|
2014-12-12 01:13:01 +00:00
|
|
|
|
2015-01-08 09:58:02 +00:00
|
|
|
def __init__(self, database, user, password, host, port):
|
2014-12-13 02:28:55 +00:00
|
|
|
(self.dbname, self.user, self.password, self.host, self.port) = \
|
|
|
|
_parse_dsn(database, default_user=user,
|
|
|
|
default_password=password, default_host=host,
|
|
|
|
default_port=port)
|
2015-01-09 09:56:22 +00:00
|
|
|
self.reconnect()
|
|
|
|
|
|
|
|
def reconnect(self):
|
2014-12-13 16:04:48 +00:00
|
|
|
self.conn = psycopg2.connect(database=self.dbname, user=self.user,
|
|
|
|
password=self.password, host=self.host, port=self.port)
|
2014-11-24 06:38:45 +00:00
|
|
|
self.conn.autocommit = True
|
2014-11-22 07:43:11 +00:00
|
|
|
|
|
|
|
def run(self, sql):
|
2014-12-10 17:03:35 +00:00
|
|
|
"""Execute the sql in the database and return the results. The results
|
|
|
|
are a list of tuples. Each tuple has 3 values (rows, headers, status).
|
|
|
|
"""
|
2014-12-08 06:30:24 +00:00
|
|
|
|
2014-12-13 21:44:46 +00:00
|
|
|
if not sql: # Empty string
|
|
|
|
return [(None, None, None)]
|
|
|
|
|
2014-12-08 06:30:24 +00:00
|
|
|
# Remove spaces, eol and semi-colons.
|
2014-11-29 06:27:18 +00:00
|
|
|
sql = sql.strip()
|
2014-12-08 06:30:24 +00:00
|
|
|
sql = sql.rstrip(';')
|
|
|
|
|
|
|
|
# Check if the command is a \c or 'use'. This is a special exception
|
|
|
|
# that cannot be offloaded to `pgspecial` lib. Because we have to
|
|
|
|
# change the database connection that we're connected to.
|
|
|
|
if sql.startswith('\c') or sql.lower().startswith('use'):
|
2015-01-05 00:29:15 +00:00
|
|
|
_logger.debug('Database change command detected.')
|
2014-12-12 04:07:54 +00:00
|
|
|
try:
|
|
|
|
dbname = sql.split()[1]
|
|
|
|
except:
|
2015-01-07 08:10:56 +00:00
|
|
|
_logger.debug('Database name missing.')
|
2014-12-12 04:07:54 +00:00
|
|
|
raise RuntimeError('Database name missing.')
|
2014-12-08 06:30:24 +00:00
|
|
|
self.conn = psycopg2.connect(database=dbname,
|
|
|
|
user=self.user, password=self.password, host=self.host,
|
|
|
|
port=self.port)
|
|
|
|
self.dbname = dbname
|
2014-12-13 20:58:37 +00:00
|
|
|
self.conn.autocommit = True
|
2015-01-05 00:29:15 +00:00
|
|
|
_logger.debug('Successfully switched to DB: %r', dbname)
|
2014-12-09 06:59:25 +00:00
|
|
|
return [(None, None, 'You are now connected to database "%s" as '
|
|
|
|
'user "%s"' % (self.dbname, self.user))]
|
2014-12-08 06:30:24 +00:00
|
|
|
|
2014-11-22 07:43:11 +00:00
|
|
|
with self.conn.cursor() as cur:
|
2014-12-08 08:43:21 +00:00
|
|
|
try:
|
2015-01-05 00:29:15 +00:00
|
|
|
_logger.debug('Trying a pgspecial command. sql: %r', sql)
|
2014-12-12 04:07:54 +00:00
|
|
|
return pgspecial.execute(cur, sql)
|
2014-12-08 08:43:21 +00:00
|
|
|
except KeyError:
|
2015-01-05 00:29:15 +00:00
|
|
|
_logger.debug('Regular sql statement. sql: %r', sql)
|
2014-11-23 23:02:05 +00:00
|
|
|
cur.execute(sql)
|
2014-11-29 04:50:03 +00:00
|
|
|
|
|
|
|
# cur.description will be None for operations that do not return
|
|
|
|
# rows.
|
|
|
|
if cur.description:
|
|
|
|
headers = [x[0] for x in cur.description]
|
2014-12-10 17:03:35 +00:00
|
|
|
return [(cur.fetchall(), headers, cur.statusmessage)]
|
2014-11-29 04:50:03 +00:00
|
|
|
else:
|
2015-01-05 00:29:15 +00:00
|
|
|
_logger.debug('No rows in result.')
|
2014-12-10 17:03:35 +00:00
|
|
|
return [(None, None, cur.statusmessage)]
|
2014-11-22 07:43:11 +00:00
|
|
|
|
2014-11-23 23:02:24 +00:00
|
|
|
def tables(self):
|
2015-01-08 10:56:28 +00:00
|
|
|
""" Returns tuple (sorted_tables, columns). Columns is a dictionary of
|
|
|
|
table name -> list of columns """
|
|
|
|
columns = defaultdict(list)
|
2014-11-23 23:02:24 +00:00
|
|
|
with self.conn.cursor() as cur:
|
|
|
|
cur.execute(self.tables_query)
|
2015-01-08 10:56:28 +00:00
|
|
|
tables = [x[0] for x in cur.fetchall()]
|
|
|
|
|
|
|
|
table_set = set(tables)
|
|
|
|
cur.execute(self.columns_query)
|
|
|
|
for table, column in cur.fetchall():
|
|
|
|
if table in table_set:
|
|
|
|
columns[table].append(column)
|
|
|
|
return tables, columns
|
2014-12-12 01:13:01 +00:00
|
|
|
|
|
|
|
def databases(self):
|
|
|
|
with self.conn.cursor() as cur:
|
|
|
|
cur.execute(self.databases_query)
|
|
|
|
return [x[0] for x in cur.fetchall()]
|