# SQLAlchemy utils

import itertools
import sqlalchemy as rdb

def get_schemata_like(engine, schemata_regexp):
    """Returns the accessible schemata in the database
    referenced by the engine.
    """
    # This needs to be a union of two separate queries,
    # because (at least in PostgreSQL) the 'schemata'
    # table only lists schemata *owned* by the user.
    # The 'tables' table lists all tables (and also
    # schemata) which are *accessible* to the user.
    # We need to query both because an empty schemata
    # without any tables would not be listed in the
    # 'tables' table. This will still miss any empty
    # schemata which aren't owned by the user.
    schemata = set()
    def add(query):
        results = query.execute().fetchall()
        if len(results) > 0:
            schemata.update(zip(*results)[0])
    # ischema.schemata
    sc = engine.ischema.schemata.columns
    query = rdb.select([sc.schema_name],
                       sc.schema_name.like(schemata_regexp))
    add(query)
    # ischema.tables
    tc = engine.ischema.tables.columns
    query = rdb.select([tc.table_schema],
                       tc.table_schema.like(schemata_regexp),
                       distinct=True)
    add(query)
    return schemata

def create_schemata_like(engine, schemata_prefix,
                         suffix_iter=None,
                         num_tries=10):
    """Creates a new, empty schema whose name starts with
    schemata_prefix. The default behavior is to try the
    suffixes '_N' where N = (0, 1, 2, ...) until the
    candidate schema name isn't already used, and the
    CREATE is successful. (The CREATE can fail if, e.g.,
    there are multiple processes attempting the same
    create_schemata_like() call.) The suffix used can
    be customized by providing and iterable for the kwarg
    'suffix_iter'. The maximum number of CREATEs attempted
    can be controlled by the kwarg 'num_tries'.
    """
    previous_schemata = get_schemata_like(engine, schemata_prefix+'%')
    if suffix_iter is None:
        suffix_iter = itertools.count()
    n = 0
    while (n < num_tries):
        suffix = str(suffix_iter.next())
        schema = schemata_prefix+'_'+suffix
        if schema in previous_schemata:
            continue
        n += 1
        # Probably better to leverage SQLAlchemy's
        # bind parameters, but don't know how to do
        # that for a custom string query.
        sql = 'CREATE SCHEMA %s;' % schema
        try:
            engine.execute(sql, [])
        except:
            # Catch everything for now.
            # Doesn't look like SQLAlchemy provides a
            # backend-independent exception to catch ...
            pass
        else:
            return schema
    raise RuntimeError("Unable to create schema like %s after %d tries. Perhaps there's a permissions problem?" % (schemata_prefix, num_tries))