def test_order_by():
     query = query_module.QueryOptions(order_by=[
         query_module.PropertyOrder("a"),
         query_module.PropertyOrder("b", reverse=True),
     ])
     expected_pb = query_pb2.Query(order=[
         query_pb2.PropertyOrder(
             property=query_pb2.PropertyReference(name="a"),
             direction=query_pb2.PropertyOrder.ASCENDING,
         ),
         query_pb2.PropertyOrder(
             property=query_pb2.PropertyReference(name="b"),
             direction=query_pb2.PropertyOrder.DESCENDING,
         ),
     ])
     assert _datastore_query._query_to_protobuf(query) == expected_pb
Ejemplo n.º 2
0
def _query_to_protobuf(query):
    """Convert an NDB query to a Datastore protocol buffer.

    Args:
        query (query.QueryOptions): The query spec.

    Returns:
        query_pb2.Query: The protocol buffer representation of the query.
    """
    query_args = {}
    if query.kind:
        query_args["kind"] = [query_pb2.KindExpression(name=query.kind)]

    if query.projection:
        query_args["projection"] = [
            query_pb2.Projection(property=query_pb2.PropertyReference(
                name=name)) for name in query.projection
        ]

    if query.distinct_on:
        query_args["distinct_on"] = [
            query_pb2.PropertyReference(name=name)
            for name in query.distinct_on
        ]

    if query.order_by:
        query_args["order"] = [
            query_pb2.PropertyOrder(
                property=query_pb2.PropertyReference(name=order.name),
                direction=DOWN if order.reverse else UP,
            ) for order in query.order_by
        ]

    filter_pb = query.filters._to_filter() if query.filters else None

    if query.ancestor:
        ancestor_pb = query.ancestor._key.to_protobuf()
        ancestor_filter_pb = query_pb2.PropertyFilter(
            property=query_pb2.PropertyReference(name="__key__"),
            op=query_pb2.PropertyFilter.HAS_ANCESTOR,
        )
        ancestor_filter_pb.value.key_value.CopyFrom(ancestor_pb)

        if filter_pb is None:
            filter_pb = ancestor_filter_pb

        elif isinstance(filter_pb, query_pb2.CompositeFilter):
            filter_pb.filters.add(property_filter=ancestor_filter_pb)

        else:
            filter_pb = query_pb2.CompositeFilter(
                op=query_pb2.CompositeFilter.AND,
                filters=[
                    _filter_pb(filter_pb),
                    _filter_pb(ancestor_filter_pb),
                ],
            )

    if filter_pb is not None:
        query_args["filter"] = _filter_pb(filter_pb)

    if query.start_cursor:
        query_args["start_cursor"] = query.start_cursor.cursor

    if query.end_cursor:
        query_args["end_cursor"] = query.end_cursor.cursor

    query_pb = query_pb2.Query(**query_args)

    if query.offset:
        query_pb.offset = query.offset

    if query.limit:
        query_pb.limit.value = query.limit

    return query_pb