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)))