from __future__ import division, print_function

# Third-party
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.ext.declarative import declarative_base

Session = scoped_session(sessionmaker(autoflush=True, autocommit=False))
Base = declarative_base()

def db_connect(host="localhost", port=5432, user=None, dbname=None, password=None):
    """
    Establish a connection to the Postgres database.
    After running :func:`connect`, sessions can be established.
    >> import starplex
    >> starplex.database.connect(**args)
    >> session = starplex.database.Session()
    Parameters
    ----------
    host : str
        Hostname
    port : int
        Server port
    user : str
        Username for server
    password : str
        Password to log into server
    name : str
        Name of database
    kwargs : dict
        Additional keyword arguments passed to ``sqlalchemy.create_engine``.
    """

    url = "postgresql://{user}:{password}@{host}:{port:d}/{dbname}".format(host=host,
                                                                           port=port,
                                                                           user=user,
                                                                           dbname=dbname,
                                                                           password=password)

    engine = create_engine(url)
    Session.configure(bind=engine)
    Base.metadata.bind = engine

    return engine

@@ -42,16 +42,46 @@
 nction%0A%0A
+# Standard library%0Aimport os%0A%0A
 # Third-
@@ -86,16 +86,16 @@
 d-party%0A
-
 from sql
@@ -234,16 +234,93 @@
 ive_base
+%0Afrom sqlalchemy.exc import OperationalError%0A%0A# Project%0Afrom ..log import log
 %0A%0ASessio
@@ -494,16 +494,54 @@
 ord=None
+,%0A               ensure_db_exists=True
 ):%0A    %22
@@ -1551,16 +1551,16 @@
 sword)%0A%0A
-
     engi
@@ -1583,16 +1583,707 @@
 ne(url)%0A
+    try:%0A        conn = engine.connect()%0A        conn.execute(%22select * from information_schema.tables%22)%0A        log.debug(%22Database '%7B%7D' exists%22.format(dbname))%0A%0A    except OperationalError:%0A        if ensure_db_exists:%0A            log.info(%22Database '%7B%7D' does not exist -- creating for you...%22.format(dbname))%0A%0A            # make sure the database exists%0A            _url = os.path.join(os.path.dirname(str(url)), 'postgres')%0A            _engine = create_engine(_url)%0A            _conn = _engine.connect()%0A            _conn.execute(%22commit%22)%0A            _conn.execute(%22CREATE DATABASE %7B0%7D%22.format(dbname))%0A            _conn.close()%0A            del _engine%0A        else:%0A            raise%0A%0A
     Sess

create db if not there