Как загрузить коллекцию с ограниченным количеством в Sqlalchemy?

Question

У меня есть две таблицы. С помощью Sqlalchemy я сопоставляю их двум классам:

class A(base):
  ...
  id = Column(BigInteger, primary_key=True, autoincrement=True)

class B(base):
  ...
  id = Column(BigInteger, primary_key=True, autoincrement=True)
  a_id = Column(BigInteger, ForeignKey(A.id))
  timestamp = Column(DateTime)

  a = relationship(A, backref="b_s")

Я могу использовать A.b_s, чтобы получить коллекцию объектов B, внешний ключ которых совпадает с первичным ключом A. Очень легко использовать ленивую или энергичную загрузку. Но теперь у меня есть вопрос. Я не хочу загружать все объекты B. Я хочу загрузить только первые N объектов, упорядоченных по метке времени. То есть A.b_s загружает только некоторые связанные объекты B. Как я могу использовать Sqlalchemy, чтобы сделать это?

Большое спасибо!

Answer 1

То, что вы хотите достичь, не будет работать с отношениями (и это не ограничение SA, а правильный способ обработки отношений и обеспечения целостности ссылок).
Однако простой запрос (заключенный в метод) отлично справится с задачей:

class A(Base):
    # ...
    def get_children(self, offset, count):
        # @todo: might need to handle some border cases
        qry = B.query.with_parent(self)
        #or: qry = object_session(self).query(B).with_parent(self)
        return qry[offset:offset+count]

my_a = session.query(A).get(a_id)
print my_a.get_children( 0, 10) # print first 10 children
print my_a.get_children(10, 10) # print second 10 children

---

edit-1: достигните этого, имея только 1-2 оператора SQL
Теперь достичь этого можно только с помощью 1-2 операторов SQL.
Прежде всего, нужен способ получить идентификаторы B для top N каждого A. Для этого мы будем использовать функцию sqlalchemy.sql.expression.over для составления подзапроса:

# @note: this is the subquery using *sqlalchemy.orm.over* function to limit number of rows
# this subquery is used for both queries below
# @note: the code below sorts Bs by id, but you can change it in order_by
subq = (session.query(
            B.__table__.c.id.label("b_id"), 
            over(func.row_number(), partition_by="a_id", order_by="id").label("rownum")
       ).subquery())
# this produces the following SQL (@note: the RDBMS should support the OVER...)
# >> SELECT b.id AS b_id, row_number() OVER (PARTITION BY a_id ORDER BY id) AS rownum FROM b

Версия-1: Теперь первая версия будет загружать A с, вторая - B с. Функция возвращает словарь с A s в качестве ключей и список B s в качестве значений:

def get_A_with_Bs_in_batch(b_limit=10):
    """ 
    @return: dict(A, [list of top *b_limit* A.b_s])  
    @note: uses 2 SQL statements, but does not screw up relationship.
    @note: if the relationship is requested via a_instance.b_s, the new SQL statement will be
    issued to load *all* related objects
    """
    qry_a = session.query(A)
    qry_b = (session.query(B)
            .join(subq, and_(subq.c.b_id == B.id, subq.c.rownum <= b_limit))
            )
    a_s = qry_a.all()
    b_s = qry_b.all()
    res = dict((a, [b for b in b_s if b.a == a]) for a in a_s)
    return res

Версия-2: это трюк SQLAlchemy, чтобы думать, что TOP N Bs, которые загружаются в одном запросе с A, на самом деле A.b_s. ОЧЕНЬ ОПАСНО, но аккуратно. Прочитайте комментарии в коде, который объясняет кусочки:

def get_A_with_Bs_hack_relation(b_limit=10):
    """ 
    @return: dict(A, [list of top *b_limit* A.b_s])
    @note: the Bs are loaded as relationship A.b_s, but with the limit.
    """
    qry = (session.query(A)
            .outerjoin(B)
            # @note: next line will trick SA to load joined Bs as if they were *all* objects
            # of relationship A.b_s. this is a @hack: and one should discard/reset a session after this
            # kind of hacky query!!!
            .options(contains_eager(A.b_s))
            .outerjoin(subq, and_(subq.c.b_id == B.id, subq.c.rownum <= b_limit))
            # @note: next line is required to make both *outerjoins* to play well together 
            # in order produce the right result
            .filter(or_(B.id == None, and_(B.id != None, subq.c.b_id != None)))
            )
    res = dict((a, a.b_s) for a in qry.all())
    return res

Подводя итог, можно сказать, что Version-2 , вероятно, является самым прямым ответом на ваш вопрос. Используйте его на свой страх и риск, потому что здесь вы обманываете SA, и если вы каким-либо образом модифицируете свойство отношений, вы можете увидеть «Kaboom!»