kbcstorage snippets

connection_tools.py

'''
Utilities for manipulating sqlalchemy connection to keboola workspace.

Example usage::

    from getpass import getpass

    from kbcstorage.client import Client
    import sqlalchemy as sqa

    from connection_tools import construct_dsn

    token = getpass('Storage token: ')
    workspace_id = input('Workspace id: ')

    client = Client("https://connection.keboola.com", token)

    engine = sqa.create_engine(construct_dsn(workspace_id, client))
    conn = engine.connect()
    q = """\\
    SELECT distinct tablename
    FROM pg_table_def
    WHERE schemaname='workspace_{}';""".format(workspace_id)
    print(conn.execute(q).fetchall())

Requires `kbcstorage`_, `sqlalchemy`_ and possibly sqlalchmey `dialects`_.

.. _kbcstorage: https://github.com/keboola/sapi-python-client
.. _slqalchemy: http://www.sqlalchemy.org/
.. _dialects: http://docs.sqlalchemy.org/en/latest/dialects/index.html
'''


def construct_dsn(workspace_id, client, ws_password=None, dialect=None,
                  driver='psycopg2', port='5439'):
    """
    Create a new sqlalchemy connection for a given workspace.

    Necessarily resets the password on creation.

    Arguments:
        workspace_id (str): The id of the workspace.
        client (kbcstorage.Client): A client instance initialised with
            a token with sufficient privileges.
        ws_password (str): The corresponding workspace password. If not
            supplied, the workspace password will be reset. Default None.
        dialect (str): the sqlalchemy dialect to use. eg. redshift. For more
            information about sqlachemy dialects, look here:
            http://docs.sqlalchemy.org/en/latest/dialects/index.html

            redshift and snowflake will require the installation of additional
            drivers.

            Default ``'redshift'``.
        driver (str): similar to dialect. May be related to dialect? Default
            'psycopg2'.
        port (str): The port to communicate with. Default 5439.
    """
    details = client.workspaces.detail(workspace_id)['connection']

    # password
    if ws_password is not None:
        details['password'] = ws_password
    else:
        details.update(client.workspaces.reset_password(workspace_id))

    # dialect
    details['dialect'] = details.pop('backend')
    if dialect is not None:
        details['dialect'] = dialect
    t = '{dialect}+{driver}://{user}:{password}@{host}:{port}/{database}'
    return t.format(driver=driver, port=port, **details)


if __name__ == "__main__":
    from getpass import getpass

    from kbcstorage.client import Client
    import sqlalchemy as sqa

    token = getpass('Storage token: ')
    workspace_id = input('Workspace id: ')

    client = Client("https://connection.keboola.com", token)

    engine = sqa.create_engine(construct_dsn(workspace_id, client))
    conn = engine.connect()
    q = """\
    SELECT distinct tablename
    FROM pg_table_def
    WHERE schemaname='workspace_{}';""".format(workspace_id)
    print(conn.execute(q).fetchall())