queries/utils.py

Summary

Maintainability
A
1 hr
Test Coverage
"""
Utility functions for access to OS level info and URI parsing

"""
import collections
import getpass
import logging
import os
import platform

# All systems do not support pwd module
try:
    import pwd
except ImportError:
    pwd = None

# Python 2 & 3 compatibility
try:
    from urllib import parse as _urlparse
except ImportError:
    import urlparse as _urlparse
try:
    from urllib.parse import unquote
except ImportError:
    from urllib import unquote

LOGGER = logging.getLogger(__name__)

PARSED = collections.namedtuple('Parsed',
                                'scheme,netloc,path,params,query,fragment,'
                                'username,password,hostname,port')

PYPY = platform.python_implementation().lower() == 'pypy'

KEYWORDS = ['connect_timeout',
            'client_encoding',
            'options',
            'application_name',
            'fallback_application_name',
            'keepalives',
            'keepalives_idle',
            'keepalives_interval',
            'keepalives_count',
            'sslmode',
            'requiressl',
            'sslcompression',
            'sslcert',
            'sslkey',
            'sslrootcert',
            'sslcrl',
            'requirepeer',
            'krbsrvname',
            'gsslib',
            'service']


def get_current_user():
    """Return the current username for the logged in user

    :rtype: str

    """
    if pwd is None:
        return getpass.getuser()
    else:
        try:
            return pwd.getpwuid(os.getuid())[0]
        except KeyError as error:
            LOGGER.error('Could not get logged-in user: %s', error)


def parse_qs(query_string):
    """Return the parsed query string in a python2/3 agnostic fashion

    :param str query_string: The URI query string
    :rtype: dict

    """
    return _urlparse.parse_qs(query_string)


def uri(host='localhost', port=5432, dbname='postgres', user='postgres',
        password=None):
    """Return a PostgreSQL connection URI for the specified values.

    :param str host: Host to connect to
    :param int port: Port to connect on
    :param str dbname: The database name
    :param str user: User to connect as
    :param str password: The password to use, None for no password
    :return str: The PostgreSQL connection URI

     """
    if port:
        host = '%s:%s' % (host, port)
    if password:
        return 'postgresql://%s:%s@%s/%s' % (user, password, host, dbname)
    return 'postgresql://%s@%s/%s' % (user, host, dbname)


def uri_to_kwargs(uri):
    """Return a URI as kwargs for connecting to PostgreSQL with psycopg2,
    applying default values for non-specified areas of the URI.

    :param str uri: The connection URI
    :rtype: dict

    """
    parsed = urlparse(uri)
    default_user = get_current_user()
    password = unquote(parsed.password) if parsed.password else None
    kwargs = {'host': parsed.hostname,
              'port': parsed.port,
              'dbname': parsed.path[1:] or default_user,
              'user': parsed.username or default_user,
              'password': password}
    values = parse_qs(parsed.query)
    if 'host' in values:
        kwargs['host'] = values['host'][0]
    for k in [k for k in values if k in KEYWORDS]:
        kwargs[k] = values[k][0] if len(values[k]) == 1 else values[k]
        try:
            if kwargs[k].isdigit():
                kwargs[k] = int(kwargs[k])
        except AttributeError:
            pass
    return kwargs


def urlparse(url):
    """Parse the URL in a Python2/3 independent fashion.

    :param str url: The URL to parse
    :rtype: Parsed

    """
    value = 'http%s' % url[5:] if url[:5] == 'postgresql' else url
    parsed = _urlparse.urlparse(value)
    path, query = parsed.path, parsed.query
    hostname = parsed.hostname if parsed.hostname else ''
    return PARSED(parsed.scheme.replace('http', 'postgresql'),
                  parsed.netloc,
                  path,
                  parsed.params,
                  query,
                  parsed.fragment,
                  parsed.username,
                  parsed.password,
                  hostname.replace('%2F', '/').replace('%2f', '/'),
                  parsed.port)