from impala.dbapi import (
    connect,
    ProgrammingError,
    DatabaseError,
    IntegrityError,
    OperationalError,
)
from impala.error import HiveServer2Error


def get_DEVO_connection(hostname: str, user: str, secret: str):
        conn = connect(
            host=hostname,
            port=443,
            auth_mechanism="PLAIN",
            user=user,
            password=secret,
            use_http_transport=True,
            http_path="cliservice",
            use_ssl=True,
        )
        return conn


def execute_devo_query(query: str, conn):
    cursor = None
    try:
        cursor = conn.cursor()
        cursor.execute(query)        # Check if the query is a SELECT query (i.e., reads data)
        if query.strip().lower().startswith("select"):
            rows = cursor.fetchall()
            columns = [col[0] for col in cursor.description]
            return columns, rows
        else:
            # For non-SELECT queries (e.g., INSERT, UPDATE, DELETE), just return affected rows
            return None, cursor.rowcount  # rowcount returns the number of rows affected

    except OperationalError as oe:
        raise Exception("Failed to connect to DEVO: " + str(oe))

    except ProgrammingError as pe:
        raise Exception("Query syntax error: " + str(pe))

    except IntegrityError as ie:
        raise Exception("Insufficient permissions: " + str(ie))

    except DatabaseError as db_err:
        raise Exception("Database error: " + str(db_err))

    except HiveServer2Error as au_err:
        raise Exception("HiveServer2Error error: " + str(au_err))

    finally:
        try:
            if cursor:
                cursor.close()
            if not conn:
                conn.close()
        except Exception as e:
            raise Exception(status_code=500, detail=f"Failed to close the cursor or impala connection: {str(e)}") from e

def execute_query(query: str, user: str, hostname: str,password):    
    conn = get_DEVO_connection(hostname, user, password)
    columns, result = execute_devo_query(query, conn)
    return columns, result


#sql="CREATE EXTERNAL TABLE IF NOT EXISTS crp_rar.testInternalTable ( iid STRING,RANDOM_DATE DATE, number int) ;"
#sql_drop="DROP TABLE IF EXISTS crp_rar.NH_PRICE"

#print( execute_query("SELECT 1","ap-informatica-ipcwt","t-impala.devo.escb.eu","Start_123456789"))
#print( execute_query("SELECT 1","ap-devo_tst-mrds","t-impala.devo.escb.eu","V1XqZ*#fvwQl=nRG*idI"))
#print( execute_query("SELECT 1","ap-devo_lab-mrds","impala-proxy-devo-lab21-impala01.dw-devo-lab21.om2y56.b0.cloudera.site","PHkvyVonyePAmZD8wUuw!"))