Fusion-SQL Access Class

fsql.py

import requests
import pandas as pd
import json
import logging


class FusionSQL:
    def __init__(self, proxy_url, user, passwd):
        self.proxy_url = proxy_url
        self.user = user
        self.passwd = passwd
        self.session = self.get_session()

    def get_session(self):
        session = requests.Session()
        logging.info(f"Connecting to {self.proxy_url}")
        resp = session.post(
            f"{self.proxy_url}/session",
            data=json.dumps({"username": self.user, "password": self.passwd}),
            headers={"Content-type": "application/json"},
        )
        if resp.status_code != 201:
            raise RuntimeError(
                f"Auth failed. status code {resp.status_code}. Response text {resp.text}"
            )
        else:
            logging.info(
                f"Auth successful for user '{self.user}' for host '{self.proxy_url}'"
            )
        return session

    def run_sql_query(self, sql):
        sql_url = f"{self.proxy_url}/catalog/fusion/query"

        sql_payload = {"sql": sql}
        try:
            resp = self.session.post(sql_url, data=json.dumps(sql_payload))
            if resp.status_code != 200:
                if resp.status_code == 401:
                    # Looks like a reconnect is needed.
                    logging.info("Attempting a reconnect...")
                    self.session = self.get_session()
                    logging.info("reconnected")
                    # Try the query again
                    resp = self.session.post(sql_url, data=json.dumps(sql_payload))
                    if resp.status_code != 200:
                        raise RuntimeError(
                            f"Failed request. Status code {resp.status_code} Response {resp.text}. Request payload {sql_payload}"
                        )
                else:
                    raise RuntimeError(
                        f"Failed request. Status code {resp.status_code} Response {resp.text}. Request payload {sql_payload}"
                    )
            else:
                logging.info(
                    f"Time taken for response {resp.elapsed.microseconds / 1000}ms"
                )
            logging.info(f"Sending sql request {sql} to sql server")
            return pd.read_json(resp.text)
        except Exception as e:
            logging.error(f"Request failed with exception {e}")
            return RuntimeError(e)