###################################################################
#           Copyright (c) 2020 by TAOS Technologies, Inc.
#                     All rights reserved.
#
#  This file is proprietary and confidential to TAOS Technologies.
#  No part of this file may be reproduced, stored, transmitted,
#  disclosed or used in any form or by any means other than as
#  expressly provided by the written permission from Jianhui Tao
#
###################################################################

# -*- coding: utf-8 -*-

import sys
import os
import time
import datetime
from util.log import *


class TDSub:
    def __init__(self):
        self.consumedRows = 0
        self.consumedCols = 0

    def init(self, sub):
        self.sub = sub

    def close(self, keepProgress):
        self.sub.close(keepProgress)

    def consume(self):
        self.result = self.sub.consume()
        self.result.fetch_all()
        self.consumedRows = self.result.row_count
        self.consumedCols = self.result.field_count
        return self.consumedRows

    def checkRows(self, expectRows):
        if self.consumedRows != expectRows:
            tdLog.exit(
                "consumed rows:%d != expect:%d" % (self.consumedRows, expectRows)
            )
        tdLog.info("consumed rows:%d == expect:%d" % (self.consumedRows, expectRows))


tdSub = TDSub()