from .cinterface import CTaosInterface
from .error import *

class TDengineSubscription(object):
    """TDengine subscription object
    """
    def __init__(self, sub):
        self._sub = sub
    

    def consume(self):
        """Consume rows of a subscription
        """
        if self._sub is None:
            raise OperationalError("Invalid use of consume")
        
        result, fields = CTaosInterface.consume(self._sub)
        buffer = [[] for i in range(len(fields))]
        while True:
            block, num_of_fields = CTaosInterface.fetchBlock(result, fields)
            if num_of_fields == 0: break
            for i in range(len(fields)):
                buffer[i].extend(block[i])

        self.fields = fields
        return list(map(tuple, zip(*buffer)))


    def close(self, keepProgress = True):
        """Close the Subscription.
        """
        if self._sub is None:
            return False
        
        CTaosInterface.unsubscribe(self._sub, keepProgress)
        return True


if __name__ == '__main__':
    from .connection import TDengineConnection
    conn = TDengineConnection(host="127.0.0.1", user="root", password="taosdata", database="test")

    # Generate a cursor object to run SQL commands
    sub = conn.subscribe(True, "test", "select * from meters;", 1000)

    for i in range(0,10):
        data = sub.consume()
        for d in data:
            print(d)

    sub.close()
    conn.close()