from urllib.parse import urlparse
import psycopg2

import sqlalchemy
from testcontainers.postgres import PostgresContainer

psql_config = {
    'database': 'test',
    'user': 'test',
    'password': 'test',
    'host': '127.0.0.1',
    'port': '5432',
}


class Cursor(object):
    def __init__(self, sql_config=psql_config, as_dict=True):
        cursor_params = dict()
        if as_dict:
            cursor_params['cursor_factory'] = psycopg2.extras.RealDictCursor
        self.connection = psycopg2.connect(**sql_config)
        self.cursor = self.connection.cursor(**cursor_params)

    def __iter__(self):
        for item in self.cursor:
            yield item

    def __enter__(self):
        return self.cursor

    def __exit__(self, ext_type, exc_value, traceback):
        self.cursor.close()
        if isinstance(exc_value, Exception):
            self.connection.rollback()
        else:
            self.connection.commit()
        self.connection.close()


def extract_config_from_url(url):
    result = urlparse(url)
    config = dict(
        user = result.username,
        password = result.password,
        database = result.path[1:],
        host = result.hostname,
        port = result.port,
    )
    return config

def test_docker_run_postgress():
    postgres_container = PostgresContainer("postgres:9.5")
    with postgres_container as postgres:
        url = postgres.get_connection_url()
        config = extract_config_from_url(url)
        with Cursor(config) as c:
            c.execute('select 5')
            print(c.fetchone())
#        e = sqlalchemy.create_engine(postgres.get_connection_url())
#_       result = e.execute("select version()")
    print('a')

test_docker_run_postgress()