SQLAlchemy中的编译扩展：索引中使用的函数

Question

我想在SQLAlchemy中创建一个不区分大小写的唯一约束，这个约束在Postgres和SQLite中都能用。

在Postgres中，可以这样实现：

• SQL: CREATE UNIQUE INDEX my_index ON my_table (lower(my_field));（创建一个唯一索引，字段值转为小写）
• SQLAlchemy: Index('my_index', func.lower(my_field), unique=True)（在SQLAlchemy中创建同样的唯一索引）

在SQLite中，可以这样实现：

• SQL: CREATE UNIQUE INDEX my_index ON my_table (my_field COLLATE NOCASE);（创建一个不区分大小写的唯一索引）
• SQLAlchemy: Index('my_index', collate(my_field, 'nocase'), unique=True)（在SQLAlchemy中创建同样的索引）

所以我想要继承FunctionElement，来创建一个自己的函数结构，这样可以根据不同的数据库系统生成my_field COLLATE NOCASE或lower(my_field)。我尝试按照编译扩展的指南进行操作，但遇到了错误（见下面的错误追踪），我能通过SQLAlchemy的源代码找到错误，但不太明白具体原因。我还尝试继承其他东西，比如ColumnElement、ColumnClause、ClauseElement，但也遇到了类似的错误。

下面是重现错误的代码。我使用的是SQLAlchemy的0.8.2版本。

from sqlalchemy import create_engine, Column, String, Integer
from sqlalchemy.types import TypeDecorator, VARCHAR
from sqlalchemy import func, Index, collate
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import (
    # Not sure which one to subclass
    FunctionElement, ColumnElement, ColumnClause, ClauseElement
)


engine = create_engine('sqlite:///:memory:', echo=True)
Base = declarative_base()


class CaseInsensitive(FunctionElement):
    name = 'CaseInsensitive'
    type = VARCHAR()


@compiles(CaseInsensitive, 'sqlite')
def case_insensitive_sqlite(element, compiler, **kw):
    arg1, = list(element.clauses)
    return collate(compiler.process(arg1), 'nocase')


@compiles(CaseInsensitive, 'postgresql')
def case_insensitive_postgresql(element, compiler, **kw):
    arg1, = list(element.clauses)
    return func.lower(compiler.process(arg1))


class MyTable(Base):

    __tablename__ = 'my_table'

    id = Column(Integer, primary_key=True)

    my_field = Column(String)

    __table_args__ = (
        Index('idx', CaseInsensitive(my_field), unique=True),
    )

---

Traceback (most recent call last):
  File "tmp.py", line 33, in <module>
    class MyTable(Base):
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/ext/declarative/api.py", line 50, in __init__
    _as_declarative(cls, classname, cls.__dict__)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/ext/declarative/base.py", line 222, in _as_declarative
    **table_kw)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/schema.py", line 332, in __new__
    table._init(name, metadata, *args, **kw)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/schema.py", line 400, in _init
    self._init_items(*args)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/schema.py", line 65, in _init_items
    item._set_parent_with_dispatch(self)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/events.py", line 236, in _set_parent_with_dispatch
    self._set_parent(parent)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/schema.py", line 2421, in _set_parent
    ColumnCollectionMixin._set_parent(self, table)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/schema.py", line 2035, in _set_parent
    self.columns.add(col)
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/sql/expression.py", line 2477, in add
    self[column.key] = column
  File "/Users/vladimir/git/wb/lib/python2.7/site-packages/sqlalchemy/sql/expression.py", line 2296, in __getattr__
    key)
AttributeError: Neither 'CaseInsensitive' object nor 'Comparator' object has an attribute 'key'

Answer 1

@compiles 方法必须返回一个字符串。此外，这里在遍历 functionelement 时，Index 有一些小问题，所以我们需要一个小小的解决办法：

class CaseInsensitive(FunctionElement):
    __visit_name__ = 'notacolumn'
    name = 'CaseInsensitive'
    type = VARCHAR()


@compiles(CaseInsensitive, 'sqlite')
def case_insensitive_sqlite(element, compiler, **kw):
    arg1, = list(element.clauses)
    return compiler.process(collate(arg1, 'nocase'), **kw)


@compiles(CaseInsensitive, 'postgresql')
def case_insensitive_postgresql(element, compiler, **kw):
    arg1, = list(element.clauses)
    return compiler.process(func.lower(arg1), **kw)