from typing import Optional
from queryservice_client.stream.segments.resultset import ResultSet
from queryservice_client.stream.resultsetstream import ResultSetStream
from queryservice_client.stream.stream import Stream

class Statement():
    def __init__(self, stream: Stream, cancel_function):
        self.stream = stream
        self.current: Optional[ResultSetStream] = None
        self.cancel_function = cancel_function

    def __iter__(self):
        return self

    def __next__(self):
        while not next(self.stream).HasField("beginResultSet"):
            pass
        self.current = ResultSetStream(self.stream)
        return ResultSet(self.current)

    def has_next(self):
        if not self.stream.has_next():
            return False
        if self.current:
            while self.current.has_next():
                next(self.current)
            self.current = None
        return self.stream.has_next()

    def cancel(self):
        self.cancel_function()