def _get_sql(query: Query):
query = str(query.compile(compile_kwargs={"literal_binds": True}))
return query
class FTS5:
def __init__(self):
self.ident = Identity.__reflection__
self.polymorph = with_polymorphic(Identity, '*')
self.sessionmaker: Callable[[], Session]
self.selectable = Query(
self.polymorph).statement.set_label_style(LS_TABLE_COL)
self.columns = self.indexed_columns()
self.rowid_c = self.translated(self.ident.mapper.c.id)
self.model_c = self.translated(self.ident.mapper.c.model)
self.idx_t = self.idx_table(self.ident.mapper)
self.idx_p = aliased(Identity, self.idx_t, adapt_on_names=True)
@property
def session(self) -> Session:
return self.sessionmaker()
@property
def initialized(self) -> bool:
return hasattr(self, 'sessionmaker')
@property
def view_name(self):
return 'identity_view'
@property
def idx_name(self):
return 'identity_idx'
def indexed_columns(self) -> List[Column]:
return [c for c in self.selectable.subquery().c]
def translated(self, target: Column) -> Column:
for col in self.selectable.subquery().c:
if col.base_columns == target.base_columns:
return col
def idx_table(self, mapper: Mapper) -> Table:
columns = []
for c in mapper.columns:
translated = self.translated(c)
args = [translated.key, c.type]
if translated.foreign_keys:
for foreign_key in translated.foreign_keys:
args.append(ForeignKey(foreign_key.column))
columns.append(Column(*args, key=c.key, primary_key=c.primary_key))
return Table(
self.idx_name,
metadata,
Column('identity_idx', types.String(), key='master'),
*columns,
keep_existing=True,
)
def polymorphic_view(self) -> DDL:
template = """
CREATE VIEW IF NOT EXISTS %(name)s
AS %(select)s
"""
info = {
'name': self.view_name,
'select': self.selectable.compile(),
}
return DDL(template % info)
def fts_virtual_table(self) -> DDL:
template = """
CREATE VIRTUAL TABLE IF NOT EXISTS %(name)s
USING fts5(%(columns)s, content=%(view_name)s, content_rowid=%(rowid_name)s)
"""
info = {
'name': self.idx_name,
'columns': ', '.join([c.key for c in self.columns]),
'view_name': self.view_name,
'rowid_name': self.rowid_c.key,
}
return DDL(template % info)
def init(self, sessionmaker: Callable[[], Session]):
self.sessionmaker = sessionmaker
session = self.session
view = self.polymorphic_view()
fts = self.fts_virtual_table()
view.execute(session.bind)
fts.execute(session.bind)
event.listen(session, 'before_flush', self.preflush_delete)
event.listen(session, 'after_flush', self.postflush_update)
session.commit()
def preflush_delete(self, session: Session, context, instances):
ids = [str(item.id) for item in [*session.dirty, *session.deleted]]
stmt = """
INSERT INTO %(name)s(%(name)s, rowid, %(columns)s)
SELECT 'delete', %(rowid_name)s, * FROM %(view_name)s
WHERE %(rowid_name)s IN (%(ids)s)
"""
info = {
'name': self.idx_name,
'columns': ', '.join([c.key for c in self.columns]),
'view_name': self.view_name,
'rowid_name': self.rowid_c.key,
'ids': ', '.join(ids),
}
session.execute(stmt % info)
def postflush_update(self, session: Session, context):
ids = [str(item.id) for item in [*session.new, *session.dirty]]
stmt = """
INSERT INTO %(name)s(rowid, %(columns)s)
SELECT %(rowid_name)s, * FROM %(view_name)s
WHERE %(rowid_name)s IN (%(ids)s)
"""
info = {
'name': self.idx_name,
'columns': ', '.join([c.key for c in self.columns]),
'view_name': self.view_name,
'rowid_name': self.rowid_c.key,
'ids': ', '.join(ids),
}
session.execute(stmt % info)
def destroy(self, session: Optional[Session] = None):
session = session or self.session
session.execute(f'DROP TABLE IF EXISTS {self.idx_name}')
session.execute(f'DROP VIEW IF EXISTS {self.view_name}')
session.commit()
def rebuild(self):
session = self.session
session.execute(
f"INSERT INTO {self.idx_name}({self.idx_name}) VALUES('rebuild');")
session.commit()
def query(self, q: Optional[str] = None) -> Query:
clause = self.idx_t.c.id.isnot(None)
if q is not None:
clause = clause & self.idx_t.c.master.op('match')(q)
return self.session.query(self.idx_p).filter(clause)
def tokenized(self, q: Optional[str] = None) -> str:
if q is None:
return None
return slugify(q, sep='* ') + '*'
def search(self, q: Optional[str] = None) -> Query:
return self.query(self.tokenized(q))
def instanceof(self, model: Type[T], q: Optional[str] = None) -> Query:
desc = [
m.entity.__name__ for m in model.__mapper__.self_and_descendants
]
targets = ' OR '.join([f'{self.model_c.key}:{d}' for d in desc])
if q is not None:
query = f'({targets}) AND {slugify(q, sep="* ")}*'
else:
query = f'({targets})'
return self.query(query)
@contextmanager
def using_mapper(self, model: Type[T]):
try:
metadata.remove(self.idx_t)
self.idx_t = self.idx_table(inspect(model))
self.idx_p = aliased(model, self.idx_t, adapt_on_names=True)
yield self.idx_p
finally:
metadata.remove(self.idx_t)
self.idx_t = self.idx_table(inspect(Identity))
self.idx_p = aliased(Identity, self.idx_t, adapt_on_names=True)
def ids(self, q: Optional[str] = None, raw_query=False) -> Query:
if not raw_query:
q = self.tokenized(q)
return self.session.query(self.idx_t.c.id).filter(
self.idx_t.c.master.op('match')(q))
def all(self, model: Type[T], q: Optional[str] = None) -> List[T]:
return self.instanceof(model, q).all()
def lookup(self, model: Type[T], q: Optional[str] = None) -> Query:
return self.session.query(model).filter(model.id.in_(self.ids(q)))
