db.py

from __future__ import annotations

import sys
import os
import datetime
import time
import threading
import requests
from requests.auth import HTTPBasicAuth


import taos
from util.sql import *
from util.cases import *
from util.dnodes import *
from util.log import *

import traceback
# from .service_manager import TdeInstance

from .config import Config
from .misc import Logging, CrashGenError, Helper
from .types import QueryResult

class DbConn:
    TYPE_NATIVE = "native-c"
    TYPE_REST =   "rest-api"
    TYPE_INVALID = "invalid"

    @classmethod
    def create(cls, connType, dbTarget):
        if connType == cls.TYPE_NATIVE:
            return DbConnNative(dbTarget)
        elif connType == cls.TYPE_REST:
            return DbConnRest(dbTarget)
        else:
            raise RuntimeError(
                "Unexpected connection type: {}".format(connType))

    @classmethod
    def createNative(cls, dbTarget) -> DbConn:
        return cls.create(cls.TYPE_NATIVE, dbTarget)

    @classmethod
    def createRest(cls, dbTarget) -> DbConn:
        return cls.create(cls.TYPE_REST, dbTarget)

    def __init__(self, dbTarget):
        self.isOpen = False
        self._type = self.TYPE_INVALID
        self._lastSql = None
        self._dbTarget = dbTarget

    def __repr__(self):
        return "[DbConn: type={}, target={}]".format(self._type, self._dbTarget)

    def getLastSql(self):
        return self._lastSql

    def open(self):
        if (self.isOpen):
            raise RuntimeError("Cannot re-open an existing DB connection")

        # below implemented by child classes
        self.openByType()

        Logging.debug("[DB] data connection opened: {}".format(self))
        self.isOpen = True

    def close(self):
        raise RuntimeError("Unexpected execution, should be overriden")

    def queryScalar(self, sql) -> int:
        return self._queryAny(sql)

    def queryString(self, sql) -> str:
        return self._queryAny(sql)

    def _queryAny(self, sql):  # actual query result as an int
        if (not self.isOpen):
            raise RuntimeError("Cannot query database until connection is open")
        nRows = self.query(sql)
        if nRows != 1:
            raise CrashGenError(
                "Unexpected result for query: {}, rows = {}".format(sql, nRows), 
                (CrashGenError.INVALID_EMPTY_RESULT if nRows==0 else CrashGenError.INVALID_MULTIPLE_RESULT)
            )
        if self.getResultRows() != 1 or self.getResultCols() != 1:
            raise RuntimeError("Unexpected result set for query: {}".format(sql))
        return self.getQueryResult()[0][0]

    def use(self, dbName):
        self.execute("use {}".format(dbName))

    def existsDatabase(self, dbName: str):
        ''' Check if a certain database exists '''
        self.query("show databases")
        dbs = [v[0] for v in self.getQueryResult()] # ref: https://stackoverflow.com/questions/643823/python-list-transformation
        # ret2 = dbName in dbs
        # print("dbs = {}, str = {}, ret2={}, type2={}".format(dbs, dbName,ret2, type(dbName)))
        return dbName in dbs # TODO: super weird type mangling seen, once here

    def existsSuperTable(self, stName):
        self.query("show stables")
        sts = [v[0] for v in self.getQueryResult()]
        return stName in sts

    def hasTables(self):
        return self.query("show tables") > 0

    def execute(self, sql):
        ''' Return the number of rows affected'''
        raise RuntimeError("Unexpected execution, should be overriden")

    def safeExecute(self, sql):
        '''Safely execute any SQL query, returning True/False upon success/failure'''
        try:
            self.execute(sql)
            return True # ignore num of results, return success
        except taos.error.Error as err:
            return False # failed, for whatever TAOS reason
        # Not possile to reach here, non-TAOS exception would have been thrown

    def query(self, sql) -> int: # return num rows returned
        ''' Return the number of rows affected'''
        raise RuntimeError("Unexpected execution, should be overriden")

    def openByType(self):
        raise RuntimeError("Unexpected execution, should be overriden")

    def getQueryResult(self) -> QueryResult :
        raise RuntimeError("Unexpected execution, should be overriden")

    def getResultRows(self):
        raise RuntimeError("Unexpected execution, should be overriden")

    def getResultCols(self):
        raise RuntimeError("Unexpected execution, should be overriden")

# Sample: curl -u root:taosdata -d "show databases" localhost:6020/rest/sql


class DbConnRest(DbConn):
    REST_PORT_INCREMENT = 11

    def __init__(self, dbTarget: DbTarget):
        super().__init__(dbTarget)
        self._type = self.TYPE_REST
        restPort = dbTarget.port + 11
        self._url = "http://{}:{}/rest/sql".format(
            dbTarget.hostAddr, dbTarget.port + self.REST_PORT_INCREMENT)
        self._result = None

    def openByType(self):  # Open connection        
        pass  # do nothing, always open

    def close(self):
        if (not self.isOpen):
            raise RuntimeError("Cannot clean up database until connection is open")
        # Do nothing for REST
        Logging.debug("[DB] REST Database connection closed")
        self.isOpen = False

    def _doSql(self, sql):
        self._lastSql = sql # remember this, last SQL attempted
        try:
            r = requests.post(self._url, 
                data = sql,
                auth = HTTPBasicAuth('root', 'taosdata'))         
        except:
            print("REST API Failure (TODO: more info here)")
            raise
        rj = r.json()
        # Sanity check for the "Json Result"
        if ('status' not in rj):
            raise RuntimeError("No status in REST response")

        if rj['status'] == 'error':  # clearly reported error
            if ('code' not in rj):  # error without code
                raise RuntimeError("REST error return without code")
            errno = rj['code']  # May need to massage this in the future
            # print("Raising programming error with REST return: {}".format(rj))
            raise taos.error.ProgrammingError(
                rj['desc'], errno)  # todo: check existance of 'desc'

        if rj['status'] != 'succ':  # better be this
            raise RuntimeError(
                "Unexpected REST return status: {}".format(
                    rj['status']))

        nRows = rj['rows'] if ('rows' in rj) else 0
        self._result = rj
        return nRows

    def execute(self, sql):
        if (not self.isOpen):
            raise RuntimeError(
                "Cannot execute database commands until connection is open")
        Logging.debug("[SQL-REST] Executing SQL: {}".format(sql))
        nRows = self._doSql(sql)
        Logging.debug(
            "[SQL-REST] Execution Result, nRows = {}, SQL = {}".format(nRows, sql))
        return nRows

    def query(self, sql):  # return rows affected
        return self.execute(sql)

    def getQueryResult(self):
        return self._result['data']

    def getResultRows(self):
        print(self._result)
        raise RuntimeError("TBD") # TODO: finish here to support -v under -c rest
        # return self._tdSql.queryRows

    def getResultCols(self):
        print(self._result)
        raise RuntimeError("TBD")

    # Duplicate code from TDMySQL, TODO: merge all this into DbConnNative


class MyTDSql:
    # Class variables
    _clsLock = threading.Lock() # class wide locking
    longestQuery = '' # type: str
    longestQueryTime = 0.0 # seconds
    lqStartTime = 0.0
    # lqEndTime = 0.0 # Not needed, as we have the two above already

    def __init__(self, hostAddr, cfgPath):
        # Make the DB connection
        self._conn = taos.connect(host=hostAddr, config=cfgPath) 
        self._cursor = self._conn.cursor()

        self.queryRows = 0
        self.queryCols = 0
        self.affectedRows = 0

    # def init(self, cursor, log=True):
    #     self.cursor = cursor
        # if (log):
        #     caller = inspect.getframeinfo(inspect.stack()[1][0])
        #     self.cursor.log(caller.filename + ".sql")

    def close(self):
        self._cursor.close() # can we double close?
        self._conn.close() # TODO: very important, cursor close does NOT close DB connection!
        self._cursor.close()

    def _execInternal(self, sql):        
        startTime = time.time() 
        # Logging.debug("Executing SQL: " + sql)
        # ret = None # TODO: use strong type here
        # try: # Let's not capture the error, and let taos.error.ProgrammingError pass through
        ret = self._cursor.execute(sql)
        # except taos.error.ProgrammingError as err:
        #     Logging.warning("Taos SQL execution error: {}, SQL: {}".format(err.msg, sql))
        #     raise CrashGenError(err.msg)
            
        # print("\nSQL success: {}".format(sql))
        queryTime =  time.time() - startTime
        # Record the query time
        cls = self.__class__
        if queryTime > (cls.longestQueryTime + 0.01) :
            with cls._clsLock:
                cls.longestQuery = sql
                cls.longestQueryTime = queryTime
                cls.lqStartTime = startTime

        # Now write to the shadow database
        if Config.isSet('use_shadow_db'):
            if sql[:11] == "INSERT INTO":
                if sql[:16] == "INSERT INTO db_0":
                    sql2 = "INSERT INTO db_s" + sql[16:]
                    self._cursor.execute(sql2)
                else:
                    raise CrashGenError("Did not find db_0 in INSERT statement: {}".format(sql))
            else: # not an insert statement
                pass

            if sql[:12] == "CREATE TABLE":
                if sql[:17] == "CREATE TABLE db_0":
                    sql2 = sql.replace('db_0', 'db_s')
                    self._cursor.execute(sql2)
                else:
                    raise CrashGenError("Did not find db_0 in CREATE TABLE statement: {}".format(sql))
            else: # not an insert statement
                pass
        
        return ret

    def query(self, sql):
        self.sql = sql
        try:
            self._execInternal(sql)
            self.queryResult = self._cursor.fetchall()
            self.queryRows = len(self.queryResult)
            self.queryCols = len(self._cursor.description)
        except Exception as e:
            # caller = inspect.getframeinfo(inspect.stack()[1][0])
            # args = (caller.filename, caller.lineno, sql, repr(e))
            # tdLog.exit("%s(%d) failed: sql:%s, %s" % args)
            raise
        return self.queryRows

    def execute(self, sql):
        self.sql = sql
        try:
            self.affectedRows = self._execInternal(sql)
        except Exception as e:
            # caller = inspect.getframeinfo(inspect.stack()[1][0])
            # args = (caller.filename, caller.lineno, sql, repr(e))
            # tdLog.exit("%s(%d) failed: sql:%s, %s" % args)
            raise
        return self.affectedRows

class DbTarget:
    def __init__(self, cfgPath, hostAddr, port):
        self.cfgPath  = cfgPath
        self.hostAddr = hostAddr
        self.port     = port
    
    def __repr__(self):
        return "[DbTarget: cfgPath={}, host={}:{}]".format(
            Helper.getFriendlyPath(self.cfgPath), self.hostAddr, self.port)

    def getEp(self):
        return "{}:{}".format(self.hostAddr, self.port)

class DbConnNative(DbConn):
    # Class variables
    _lock = threading.Lock()
    # _connInfoDisplayed = False # TODO: find another way to display this
    totalConnections = 0 # Not private
    totalRequests = 0 

    def __init__(self, dbTarget):
        super().__init__(dbTarget)
        self._type = self.TYPE_NATIVE
        self._conn = None
        # self._cursor = None   
        
    @classmethod
    def resetTotalRequests(cls):        
        with cls._lock: # force single threading for opening DB connections. # TODO: whaaat??!!!            
            cls.totalRequests = 0

    def openByType(self):  # Open connection
        # global gContainer
        # tInst = tInst or gContainer.defTdeInstance # set up in ClientManager, type: TdeInstance
        # cfgPath = self.getBuildPath() + "/test/cfg"
        # cfgPath  = tInst.getCfgDir()
        # hostAddr = tInst.getHostAddr()

        cls = self.__class__ # Get the class, to access class variables
        with cls._lock: # force single threading for opening DB connections. # TODO: whaaat??!!!
            dbTarget = self._dbTarget
            # if not cls._connInfoDisplayed:
            #     cls._connInfoDisplayed = True # updating CLASS variable
            Logging.debug("Initiating TAOS native connection to {}".format(dbTarget))                    
            # Make the connection         
            # self._conn = taos.connect(host=hostAddr, config=cfgPath)  # TODO: make configurable
            # self._cursor = self._conn.cursor()
            # Record the count in the class
            self._tdSql = MyTDSql(dbTarget.hostAddr, dbTarget.cfgPath) # making DB connection
            cls.totalConnections += 1 
        
        self._tdSql.execute('reset query cache')
        # self._cursor.execute('use db') # do this at the beginning of every

        # Open connection
        # self._tdSql = MyTDSql()
        # self._tdSql.init(self._cursor)
        
    def close(self):
        if (not self.isOpen):
            raise RuntimeError("Cannot clean up database until connection is open")
        self._tdSql.close()
        # Decrement the class wide counter
        cls = self.__class__ # Get the class, to access class variables
        with cls._lock:
            cls.totalConnections -= 1

        Logging.debug("[DB] Database connection closed")
        self.isOpen = False

    def execute(self, sql):
        if (not self.isOpen):
            traceback.print_stack()
            raise CrashGenError(
                "Cannot exec SQL unless db connection is open", CrashGenError.DB_CONNECTION_NOT_OPEN)
        Logging.debug("[SQL] Executing SQL: {}".format(sql))
        self._lastSql = sql
        nRows = self._tdSql.execute(sql)
        cls = self.__class__
        cls.totalRequests += 1
        Logging.debug(
            "[SQL] Execution Result, nRows = {}, SQL = {}".format(
                nRows, sql))
        return nRows

    def query(self, sql):  # return rows affected
        if (not self.isOpen):
            traceback.print_stack()
            raise CrashGenError(
                "Cannot query database until connection is open, restarting?", CrashGenError.DB_CONNECTION_NOT_OPEN)
        Logging.debug("[SQL] Executing SQL: {}".format(sql))
        self._lastSql = sql
        nRows = self._tdSql.query(sql)
        cls = self.__class__
        cls.totalRequests += 1
        Logging.debug(
            "[SQL] Query Result, nRows = {}, SQL = {}".format(
                nRows, sql))
        return nRows
        # results are in: return self._tdSql.queryResult

    def getQueryResult(self):
        return self._tdSql.queryResult

    def getResultRows(self):
        return self._tdSql.queryRows

    def getResultCols(self):
        return self._tdSql.queryCols


class DbManager():
    ''' This is a wrapper around DbConn(), to make it easier to use. 

        TODO: rename this to DbConnManager
    '''
    def __init__(self, cType, dbTarget):
        # self.tableNumQueue = LinearQueue() # TODO: delete?
        # self.openDbServerConnection()
        self._dbConn = DbConn.createNative(dbTarget) if (
            cType == 'native') else DbConn.createRest(dbTarget)
        try:
            self._dbConn.open()  # may throw taos.error.ProgrammingError: disconnected
            Logging.debug("DbManager opened DB connection...")
        except taos.error.ProgrammingError as err:
            # print("Error type: {}, msg: {}, value: {}".format(type(err), err.msg, err))
            if (err.msg == 'client disconnected'):  # cannot open DB connection
                print(
                    "Cannot establish DB connection, please re-run script without parameter, and follow the instructions.")
                sys.exit(2)
            else:
                print("Failed to connect to DB, errno = {}, msg: {}"
                    .format(Helper.convertErrno(err.errno), err.msg))
                raise
        except BaseException:
            print("[=] Unexpected exception")
            raise

        # Do this after dbConn is in proper shape
        # Moved to Database()
        # self._stateMachine = StateMechine(self._dbConn)

    def __del__(self):
        ''' Release the underlying DB connection upon deletion of DbManager '''
        self.cleanUp()

    def getDbConn(self) -> DbConn :
        if self._dbConn is None:
            raise CrashGenError("Unexpected empty DbConn")
        return self._dbConn

    def cleanUp(self):
        if self._dbConn:
            self._dbConn.close()
            self._dbConn = None
            Logging.debug("DbManager closed DB connection...")