flexget.utils.sqlalchemy_utils
Covered: 56 lines
Missed: 28 lines
Skipped 23 lines
Percent: 66 %
  1
"""
  2
Miscellaneous SQLAlchemy helpers.
  3
"""
  4
from sqlalchemy import ColumnDefault, Sequence
  5
from sqlalchemy.types import AbstractType
  6
from sqlalchemy.schema import Table, MetaData
  7
from sqlalchemy.exc import NoSuchTableError
 10
def table_exists(name, session):
 11
    """
 12
    Use SQLAlchemy reflect to check table existences.
 14
    :param string name: Table name to check
 15
    :param Session session: Session to use
 16
    :return: True if table exists, False otherwise
 17
    :rtype: bool
 18
    """
 19
    try:
 20
        meta = MetaData(bind=session.connection())
 21
        Table(name, meta, autoload=True, autoload_with=session.connection())
 22
    except NoSuchTableError:
 23
        return False
 24
    return True
 27
def table_schema(name, session):
 28
    """
 29
    :returns: Table schema using SQLAlchemy reflect as it currently exists in the db
 30
    :rtype: Table
 31
    """
 32
    meta = MetaData(bind=session.bind, reflect=True)
 33
    for table in meta.sorted_tables:
 34
        if table.name == name:
 35
            return table
 38
def table_columns(table, session):
 39
    """
 40
    :param string table: Name of table or table schema
 41
    :param Session session: SQLAlchemy Session
 42
    :returns: List of column names in the table or empty list
 43
    """
 45
    res = []
 46
    if isinstance(table, basestring):
 47
        table = table_schema(table, session)
 48
    for column in table.columns:
 49
        res.append(column.name)
 50
    return res
 53
def table_add_column(table, name, col_type, session, default=None):
 54
    """Adds a column to a table
 56
    .. warning:: Uses raw statements, probably needs to be changed in
 57
                 order to work on other databases besides SQLite
 59
    :param string table: Table to add column to (can be name or schema)
 60
    :param string name: Name of new column to add
 61
    :param col_type: The sqlalchemy column type to add
 62
    :param Session session: SQLAlchemy Session to do the alteration
 63
    :param default: Default value for the created column (optional)
 64
    """
 65
    if isinstance(table, basestring):
 66
        table = table_schema(table, session)
 67
    if name in table_columns(table, session):
 69
        return
 71
    if not isinstance(col_type, AbstractType):
 73
        col_type = col_type()
 74
    type_string = session.bind.engine.dialect.type_compiler.process(col_type)
 75
    statement = 'ALTER TABLE %s ADD %s %s' % (table.name, name, type_string)
 76
    session.execute(statement)
 78
    if default is not None:
 80
        table = table_schema(table.name, session)
 81
        if not isinstance(default, (ColumnDefault, Sequence)):
 82
            default = ColumnDefault(default)
 83
        default._set_parent(getattr(table.c, name))
 84
        statement = table.update().values({name: default.execute(bind=session.bind)})
 85
        session.execute(statement)
 88
def drop_tables(names, session):
 89
    """Takes a list of table names and drops them from the database if they exist."""
 90
    metadata = MetaData(bind=session.bind, reflect=True)
 91
    for table in metadata.sorted_tables:
 92
        if table.name in names:
 93
            table.drop()
 96
def get_index_by_name(table, name):
 97
    """
 98
    Find declaratively defined index from table by name
100
    :param table: Table object
101
    :param string name: Name of the index to get
102
    :return: Index object
103
    """
104
    for index in table.indexes:
105
        if index.name == name:
106
            return index