from queryservice_client.protobuf.rdbms.query_pb2 import QueryResultObject

from queryservice_client.stream.stream import Stream
from queryservice_client.stream.delimitedstream import DelimitedStream
from queryservice_client.stream.peekable import Peekable

class ResultSetStream(Peekable):

    @staticmethod
    def check(obj: QueryResultObject):
        return obj.HasField("endOfResultSet")

    def __init__(self, stream: Stream):
        self.stream = DelimitedStream(stream, ResultSetStream.check)

    def has_next(self):
        return self.stream.has_next()

    def __next__(self):
        return next(self.stream)

    def peek(self):
        return self.stream.peek()

    def __iter__(self):
        return self