datapilot-cli/src/datapilot/clients/altimate/client.py at 510ec2c7389342d90e3dcad9c79467490fa22faf · AltimateAI/datapilot-cli · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import logging

import requests
from requests.exceptions import ConnectionError
from requests.exceptions import HTTPError
from requests.exceptions import RequestException
from requests.exceptions import Timeout


class APIClient:
    def __init__(self, api_token="", base_url="", tenant=""):
        self.api_token = api_token
        self.base_url = base_url
        self.tenant = tenant
        self.logger = logging.getLogger(self.__class__.__name__)

    def _get_headers(self):
        headers = {
            "Content-Type": "application/json",
        }

        if self.api_token:
            headers["Authorization"] = f"Bearer {self.api_token}"

        if self.tenant:
            headers["x-tenant"] = self.tenant

        return headers

    def log(self, message):
        self.logger.debug(message)

    def get(self, endpoint, params=None, timeout=None):
        url = f"{self.base_url}{endpoint}"
        headers = self._get_headers()

        try:
            self.logger.debug(f"Sending GET request for tenant {self.tenant} at url: {url}")
            response = requests.get(url, headers=headers, params=params, timeout=timeout)

            # Check if the response was successful
            response.raise_for_status()
            self.logger.debug(f"Received GET response with status: {response.status_code}")
            return response.json()

        except HTTPError as http_err:
            self.logger.error(f"{http_err.response.json()} - Status code: {response.status_code}")
        except ConnectionError as conn_err:
            self.logger.error(f"Connection error occurred: {conn_err}")
        except Timeout as timeout_err:
            self.logger.error(f"Timeout error occurred: {timeout_err}")
        except RequestException as req_err:
            self.logger.error(f"Unexpected error occurred: {req_err}")
        except Exception as err:
            self.logger.error(f"An error occurred: {err}")

    def post(self, endpoint, data=None, timeout=None):
        url = f"{self.base_url}{endpoint}"
        headers = self._get_headers()

        self.logger.debug(f"Sending POST request for tenant {self.tenant} at url: {url}")
        response = requests.post(url, headers=headers, json=data, timeout=timeout)
        self.logger.debug(f"Received POST response with status: {response.status_code }")

        return response

    def put(self, endpoint, data, timeout=None):
        url = f"{self.base_url}{endpoint}"

        self.logger.debug(f"Sending PUT request for tenant {self.tenant} at url: {url}")
        response = requests.put(url, data=data, timeout=timeout)
        self.logger.debug(f"Received PUT response with status: {response.status_code}")
        return response

    def verify_upload(self, params=None):
        endpoint = "/dbt/v1/verify_upload"
        self.post(endpoint, data=params)

    def get_signed_url(self, params=None):
        endpoint = "/dbt/v1/signed_url"
        return self.get(endpoint, params=params)

    def validate_credentials(self):
        endpoint = "/dbt/v3/validate-credentials"
        return self.get(endpoint)

    def sql_insights(self, data):
        endpoint = "/dbt/v1/sqlinsights"
        return self.post(endpoint, data=data)