Python sqlalchemy.ext.compiler.compiles() Examples

def test_functions(self):
        from sqlalchemy.dialects import postgresql

        class MyUtcFunction(FunctionElement):
            pass

        @compiles(MyUtcFunction)
        def visit_myfunc(element, compiler, **kw):
            return "utcnow()"

        @compiles(MyUtcFunction, "postgresql")
        def visit_myfunc_pg(element, compiler, **kw):
            return "timezone('utc', current_timestamp)"

        self.assert_compile(
            MyUtcFunction(), "utcnow()", use_default_dialect=True
        )
        self.assert_compile(
            MyUtcFunction(),
            "timezone('utc', current_timestamp)",
            dialect=postgresql.dialect(),
        ) 

def test_truly_unlabeled_sql_expressions(self):
        users = self.tables.users
        sess = Session()

        class not_named_max(expression.ColumnElement):
            name = "not_named_max"

        @compiles(not_named_max)
        def visit_max(element, compiler, **kw):
            return "max(id)"

        # assert that there is no "AS max_" or any label of any kind.
        eq_(str(select([not_named_max()])), "SELECT max(id)")

        # ColumnElement still handles it by applying label()
        q = sess.query(not_named_max()).select_from(users)
        eq_(q.all(), [(10,)]) 

def test_column(self):
        class MyThingy(ColumnClause):
            def __init__(self, arg=None):
                super(MyThingy, self).__init__(arg or "MYTHINGY!")

        @compiles(MyThingy)
        def visit_thingy(thingy, compiler, **kw):
            return ">>%s<<" % thingy.name

        self.assert_compile(
            select([column("foo"), MyThingy()]), "SELECT foo, >>MYTHINGY!<<"
        )

        self.assert_compile(
            select([MyThingy("x"), MyThingy("y")]).where(MyThingy() == 5),
            "SELECT >>x<<, >>y<< WHERE >>MYTHINGY!<< = :MYTHINGY!_1",
        ) 

def test_create_column_skip(self):
        @compiles(CreateColumn)
        def skip_xmin(element, compiler, **kw):
            if element.element.name == "xmin":
                return None
            else:
                return compiler.visit_create_column(element, **kw)

        t = Table(
            "t",
            MetaData(),
            Column("a", Integer),
            Column("xmin", Integer),
            Column("c", Integer),
        )

        self.assert_compile(
            CreateTable(t), "CREATE TABLE t (a INTEGER, c INTEGER)"
        ) 

def test_types(self):
        class MyType(TypeEngine):
            pass

        @compiles(MyType, "sqlite")
        def visit_sqlite_type(type_, compiler, **kw):
            return "SQLITE_FOO"

        @compiles(MyType, "postgresql")
        def visit_pg_type(type_, compiler, **kw):
            return "POSTGRES_FOO"

        from sqlalchemy.dialects.sqlite import base as sqlite
        from sqlalchemy.dialects.postgresql import base as postgresql

        self.assert_compile(MyType(), "SQLITE_FOO", dialect=sqlite.dialect())

        self.assert_compile(
            MyType(), "POSTGRES_FOO", dialect=postgresql.dialect()
        ) 

def test_callout_to_compiler(self):
        class InsertFromSelect(ClauseElement):
            def __init__(self, table, select):
                self.table = table
                self.select = select

        @compiles(InsertFromSelect)
        def visit_insert_from_select(element, compiler, **kw):
            return "INSERT INTO %s (%s)" % (
                compiler.process(element.table, asfrom=True),
                compiler.process(element.select),
            )

        t1 = table("mytable", column("x"), column("y"), column("z"))
        self.assert_compile(
            InsertFromSelect(t1, select([t1]).where(t1.c.x > 5)),
            "INSERT INTO mytable (SELECT mytable.x, mytable.y, mytable.z "
            "FROM mytable WHERE mytable.x > :x_1)",
        ) 

def test_conn_execute(self, connection):
        from sqlalchemy.sql.expression import FunctionElement
        from sqlalchemy.ext.compiler import compiles

        class myfunc(FunctionElement):
            type = Date()

        @compiles(myfunc)
        def compile_(elem, compiler, **kw):
            return compiler.process(func.current_date())

        x = connection.execute(func.current_date()).scalar()
        y = connection.execute(func.current_date().select()).scalar()
        z = connection.scalar(func.current_date())
        q = connection.scalar(myfunc())

        assert (x == y == z == q) is True 

def define_tables(cls, metadata):
        from sqlalchemy.sql import ColumnElement
        from sqlalchemy.ext.compiler import compiles

        counter = itertools.count()

        class IncDefault(ColumnElement):
            pass

        @compiles(IncDefault)
        def compile_(element, compiler, **kw):
            return str(next(counter))

        Table(
            "t1",
            metadata,
            Column(
                "id", Integer, primary_key=True, test_needs_autoincrement=True
            ),
            Column("data", String(50)),
            Column("insdef", Integer, default=IncDefault()),
            Column("upddef", Integer, onupdate=IncDefault()),
        ) 

def _fixture(self):
        class SomeColThing(WrapsColumnExpression, ColumnElement):
            def __init__(self, expression):
                self.clause = coercions.expect(
                    roles.ExpressionElementRole, expression
                )

            @property
            def wrapped_column_expression(self):
                return self.clause

        @compiles(SomeColThing)
        def process(element, compiler, **kw):
            return "SOME_COL_THING(%s)" % compiler.process(
                element.clause, **kw
            )

        return SomeColThing 

def test_function_subclasses_one(self):
        class Base(FunctionElement):
            name = "base"

        class Sub1(Base):
            name = "sub1"

        class Sub2(Base):
            name = "sub2"

        @compiles(Base)
        def visit_base(element, compiler, **kw):
            return element.name

        @compiles(Sub1)
        def visit_sub1(element, compiler, **kw):
            return "FOO" + element.name

        self.assert_compile(
            select([Sub1(), Sub2()]),
            "SELECT FOOsub1 AS sub1_1, sub2 AS sub2_1",
            use_default_dialect=True,
        ) 

def test_result_map_population_explicit(self):
        class not_named_max(ColumnElement):
            name = "not_named_max"

        @compiles(not_named_max)
        def visit_max(element, compiler, **kw):
            # explicit add
            kw["add_to_result_map"](None, None, (element,), NULLTYPE)
            return "max(a)"

        nnm = not_named_max()
        self._test_result_map_population(nnm) 

def test_function_calls_base(self):
        from sqlalchemy.dialects import mssql

        class greatest(FunctionElement):
            type = Numeric()
            name = "greatest"

        @compiles(greatest)
        def default_greatest(element, compiler, **kw):
            return compiler.visit_function(element)

        @compiles(greatest, "mssql")
        def case_greatest(element, compiler, **kw):
            arg1, arg2 = list(element.clauses)
            return "CASE WHEN %s > %s THEN %s ELSE %s END" % (
                compiler.process(arg1),
                compiler.process(arg2),
                compiler.process(arg1),
                compiler.process(arg2),
            )

        self.assert_compile(
            greatest("a", "b"),
            "greatest(:greatest_1, :greatest_2)",
            use_default_dialect=True,
        )
        self.assert_compile(
            greatest("a", "b"),
            "CASE WHEN :greatest_1 > :greatest_2 "
            "THEN :greatest_1 ELSE :greatest_2 END",
            dialect=mssql.dialect(),
        ) 

def test_functions_args_noname(self):
        class myfunc(FunctionElement):
            pass

        @compiles(myfunc)
        def visit_myfunc(element, compiler, **kw):
            return "myfunc%s" % (compiler.process(element.clause_expr, **kw),)

        self.assert_compile(myfunc(), "myfunc()")

        self.assert_compile(myfunc(column("x"), column("y")), "myfunc(x, y)") 

def test_custom_create(self):
        from sqlalchemy.ext.compiler import compiles, deregister

        @compiles(schema.CreateColumn)
        def compile_(element, compiler, **kw):
            column = element.element

            if "special" not in column.info:
                return compiler.visit_create_column(element, **kw)

            text = "%s SPECIAL DIRECTIVE %s" % (
                column.name,
                compiler.type_compiler.process(column.type),
            )
            default = compiler.get_column_default_string(column)
            if default is not None:
                text += " DEFAULT " + default

            if not column.nullable:
                text += " NOT NULL"

            if column.constraints:
                text += " ".join(
                    compiler.process(const) for const in column.constraints
                )
            return text

        t = Table(
            "mytable",
            MetaData(),
            Column("x", Integer, info={"special": True}, primary_key=True),
            Column("y", String(50)),
            Column("z", String(20), info={"special": True}),
        )

        self.assert_compile(
            schema.CreateTable(t),
            "CREATE TABLE mytable (x SPECIAL DIRECTIVE INTEGER "
            "NOT NULL, y VARCHAR(50), "
            "z SPECIAL DIRECTIVE VARCHAR(20), PRIMARY KEY (x))",
        )

        deregister(schema.CreateColumn) 

def test_default_subclass(self):
        from sqlalchemy.types import ARRAY

        class MyArray(ARRAY):
            pass

        @compiles(MyArray, "sqlite")
        def sl_array(elem, compiler, **kw):
            return "array"

        self.assert_compile(
            MyArray(Integer), "INTEGER[]", dialect="postgresql"
        ) 

def test_no_default_message(self):
        class MyThingy(ClauseElement):
            pass

        @compiles(MyThingy, "postgresql")
        def visit_thingy(thingy, compiler, **kw):
            return "mythingy"

        assert_raises_message(
            exc.CompileError,
            "<class 'test.ext.test_compiler..*MyThingy'> "
            "construct has no default compilation handler.",
            str,
            MyThingy(),
        ) 

def test_result_map_population_implicit(self):
        class not_named_max(ColumnElement):
            name = "not_named_max"

        @compiles(not_named_max)
        def visit_max(element, compiler, **kw):
            # we don't add to keymap here; compiler should be doing it
            return "max(a)"

        nnm = not_named_max()
        self._test_result_map_population(nnm) 

def test_select(self):
        t1 = table("t1", column("c1"), column("c2"))

        @compiles(Select, "sqlite")
        def compile_(element, compiler, **kw):
            return "OVERRIDE"

        s1 = select([t1])
        self.assert_compile(s1, "SELECT t1.c1, t1.c2 FROM t1")

        from sqlalchemy.dialects.sqlite import base as sqlite

        self.assert_compile(s1, "OVERRIDE", dialect=sqlite.dialect()) 

def test_no_default_has_no_visit(self):
        class MyThingy(TypeEngine):
            pass

        @compiles(MyThingy, "postgresql")
        def visit_thingy(thingy, compiler, **kw):
            return "mythingy"

        assert_raises_message(
            exc.CompileError,
            "<class 'test.ext.test_compiler..*MyThingy'> "
            "construct has no default compilation handler.",
            str,
            MyThingy(),
        ) 

def test_binds_in_select(self):
        t = table("t", column("a"), column("b"), column("c"))

        @compiles(BindParameter)
        def gen_bind(element, compiler, **kw):
            return "BIND(%s)" % compiler.visit_bindparam(element, **kw)

        self.assert_compile(
            t.select().where(t.c.c == 5),
            "SELECT t.a, t.b, t.c FROM t WHERE t.c = BIND(:c_1)",
            use_default_dialect=True,
        ) 

def test_no_default_but_has_a_visit(self):
        class MyThingy(ColumnClause):
            pass

        @compiles(MyThingy, "postgresql")
        def visit_thingy(thingy, compiler, **kw):
            return "mythingy"

        eq_(str(MyThingy("x")), "x") 

def test_binds_in_dml(self):
        t = table("t", column("a"), column("b"), column("c"))

        @compiles(BindParameter)
        def gen_bind(element, compiler, **kw):
            return "BIND(%s)" % compiler.visit_bindparam(element, **kw)

        self.assert_compile(
            t.insert(),
            "INSERT INTO t (a, b) VALUES (BIND(:a), BIND(:b))",
            {"a": 1, "b": 2},
            use_default_dialect=True,
        ) 

def compiles(class_, *specs):
    """Register a function as a compiler for a
    given :class:`.ClauseElement` type."""

    def decorate(fn):
        # get an existing @compiles handler
        existing = class_.__dict__.get('_compiler_dispatcher', None)

        # get the original handler.  All ClauseElement classes have one
        # of these, but some TypeEngine classes will not.
        existing_dispatch = getattr(class_, '_compiler_dispatch', None)

        if not existing:
            existing = _dispatcher()

            if existing_dispatch:
                def _wrap_existing_dispatch(element, compiler, **kw):
                    try:
                        return existing_dispatch(element, compiler, **kw)
                    except exc.UnsupportedCompilationError:
                        raise exc.CompileError(
                            "%s construct has no default "
                            "compilation handler." % type(element))
                existing.specs['default'] = _wrap_existing_dispatch

            # TODO: why is the lambda needed ?
            setattr(class_, '_compiler_dispatch',
                    lambda *arg, **kw: existing(*arg, **kw))
            setattr(class_, '_compiler_dispatcher', existing)

        if specs:
            for s in specs:
                existing.specs[s] = fn

        else:
            existing.specs['default'] = fn
        return fn
    return decorate 

def compiles(class_, *specs):
    """Register a function as a compiler for a
    given :class:`.ClauseElement` type."""

    def decorate(fn):
        existing = class_.__dict__.get('_compiler_dispatcher', None)
        existing_dispatch = class_.__dict__.get('_compiler_dispatch')
        if not existing:
            existing = _dispatcher()

            if existing_dispatch:
                existing.specs['default'] = existing_dispatch

            # TODO: why is the lambda needed ?
            setattr(class_, '_compiler_dispatch',
                lambda *arg, **kw: existing(*arg, **kw))
            setattr(class_, '_compiler_dispatcher', existing)

        if specs:
            for s in specs:
                existing.specs[s] = fn

        else:
            existing.specs['default'] = fn
        return fn
    return decorate 

def compiles(class_, *specs):
    """Register a function as a compiler for a
    given :class:`.ClauseElement` type."""

    def decorate(fn):
        # get an existing @compiles handler
        existing = class_.__dict__.get('_compiler_dispatcher', None)

        # get the original handler.  All ClauseElement classes have one
        # of these, but some TypeEngine classes will not.
        existing_dispatch = getattr(class_, '_compiler_dispatch', None)

        if not existing:
            existing = _dispatcher()

            if existing_dispatch:
                def _wrap_existing_dispatch(element, compiler, **kw):
                    try:
                        return existing_dispatch(element, compiler, **kw)
                    except exc.UnsupportedCompilationError:
                        raise exc.CompileError(
                            "%s construct has no default "
                            "compilation handler." % type(element))
                existing.specs['default'] = _wrap_existing_dispatch

            # TODO: why is the lambda needed ?
            setattr(class_, '_compiler_dispatch',
                    lambda *arg, **kw: existing(*arg, **kw))
            setattr(class_, '_compiler_dispatcher', existing)

        if specs:
            for s in specs:
                existing.specs[s] = fn

        else:
            existing.specs['default'] = fn
        return fn
    return decorate 

def use_simple_names():
    get_col_name = lambda el, *args, **kwargs: str(el.element.name)
    try:
        yield compiles(sql.compiler._CompileLabel)(get_col_name)
    except:
        pass
    finally:
        deregister(sql.compiler._CompileLabel) 

def _is_sqlalchemy_connectable(con):
    global _SQLALCHEMY_INSTALLED
    if _SQLALCHEMY_INSTALLED is None:
        try:
            import sqlalchemy
            _SQLALCHEMY_INSTALLED = True

            from distutils.version import LooseVersion
            ver = LooseVersion(sqlalchemy.__version__)
            # For sqlalchemy versions < 0.8.2, the BIGINT type is recognized
            # for a sqlite engine, which results in a warning when trying to
            # read/write a DataFrame with int64 values. (GH7433)
            if ver < '0.8.2':
                from sqlalchemy import BigInteger
                from sqlalchemy.ext.compiler import compiles

                @compiles(BigInteger, 'sqlite')
                def compile_big_int_sqlite(type_, compiler, **kw):
                    return 'INTEGER'
        except ImportError:
            _SQLALCHEMY_INSTALLED = False

    if _SQLALCHEMY_INSTALLED:
        import sqlalchemy
        return isinstance(con, sqlalchemy.engine.Connectable)
    else:
        return False 

def compiles(class_, *specs):
    """Register a function as a compiler for a
    given :class:`.ClauseElement` type."""

    def decorate(fn):
        existing = class_.__dict__.get('_compiler_dispatcher', None)
        existing_dispatch = class_.__dict__.get('_compiler_dispatch')
        if not existing:
            existing = _dispatcher()

            if existing_dispatch:
                existing.specs['default'] = existing_dispatch

            # TODO: why is the lambda needed ?
            setattr(class_, '_compiler_dispatch',
                    lambda *arg, **kw: existing(*arg, **kw))
            setattr(class_, '_compiler_dispatcher', existing)

        if specs:
            for s in specs:
                existing.specs[s] = fn

        else:
            existing.specs['default'] = fn
        return fn
    return decorate 

def _is_sqlalchemy_connectable(con):
    global _SQLALCHEMY_INSTALLED
    if _SQLALCHEMY_INSTALLED is None:
        try:
            import sqlalchemy
            _SQLALCHEMY_INSTALLED = True

            from distutils.version import LooseVersion
            ver = sqlalchemy.__version__
            # For sqlalchemy versions < 0.8.2, the BIGINT type is recognized
            # for a sqlite engine, which results in a warning when trying to
            # read/write a DataFrame with int64 values. (GH7433)
            if LooseVersion(ver) < LooseVersion('0.8.2'):
                from sqlalchemy import BigInteger
                from sqlalchemy.ext.compiler import compiles

                @compiles(BigInteger, 'sqlite')
                def compile_big_int_sqlite(type_, compiler, **kw):
                    return 'INTEGER'
        except ImportError:
            _SQLALCHEMY_INSTALLED = False

    if _SQLALCHEMY_INSTALLED:
        import sqlalchemy
        return isinstance(con, sqlalchemy.engine.Connectable)
    else:
        return False 

def visit_delete_limit(element, compiler, **kw):
    """
    Default compiler for the DeleteLimit clause element.
    This compiles to a DELETE statement with a SELECT subquery which
    has a limit set::

        DELETE FROM ... WHERE id IN
        (SELECT id FROM ... WHERE ... LIMIT ...)

    However, this syntax is not supported by MySQL.
    """
    select_stmt = select([element.table.c.id]).where(element.filter).limit(element.limit)
    delete_stmt = element.table.delete().where(element.table.c.id.in_(select_stmt))
    return compiler.process(delete_stmt)