forked from databricks/databricks-sql-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathhttp.py
More file actions
152 lines (126 loc) · 4.93 KB
/
Copy pathhttp.py
File metadata and controls
152 lines (126 loc) · 4.93 KB
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from enum import Enum
import threading
from dataclasses import dataclass
from contextlib import contextmanager
from typing import Generator, Optional
import logging
from requests.adapters import HTTPAdapter
from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
logger = logging.getLogger(__name__)
# Enums for HTTP Methods
class HttpMethod(str, Enum):
GET = "GET"
POST = "POST"
PUT = "PUT"
DELETE = "DELETE"
# HTTP request headers
class HttpHeader(str, Enum):
CONTENT_TYPE = "Content-Type"
AUTHORIZATION = "Authorization"
# Dataclass for OAuthHTTP Response
@dataclass
class OAuthResponse:
token_type: str = ""
expires_in: int = 0
ext_expires_in: int = 0
expires_on: int = 0
not_before: int = 0
resource: str = ""
access_token: str = ""
refresh_token: str = ""
# Singleton class for common Http Client
class DatabricksHttpClient:
## TODO: Unify all the http clients in the PySQL Connector
_instance = None
_lock = threading.Lock()
def __init__(self):
self.session = requests.Session()
adapter = HTTPAdapter(
pool_connections=5,
pool_maxsize=10,
max_retries=Retry(total=10, backoff_factor=0.1),
)
self.session.mount("https://", adapter)
self.session.mount("http://", adapter)
@classmethod
def get_instance(cls) -> "DatabricksHttpClient":
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = DatabricksHttpClient()
return cls._instance
@contextmanager
def execute(
self, method: HttpMethod, url: str, **kwargs
) -> Generator[requests.Response, None, None]:
logger.info("Executing HTTP request: %s with url: %s", method.value, url)
response = None
try:
response = self.session.request(method.value, url, **kwargs)
yield response
except Exception as e:
logger.error("Error executing HTTP request in DatabricksHttpClient: %s", e)
raise e
finally:
if response is not None:
response.close()
def close(self):
self.session.close()
class TelemetryHTTPAdapter(HTTPAdapter):
"""
Custom HTTP adapter to prepare our DatabricksRetryPolicy before each request.
This ensures the retry timer is started and the command type is set correctly,
allowing the policy to manage its state for the duration of the request retries.
"""
def send(self, request, **kwargs):
self.max_retries.command_type = CommandType.OTHER
self.max_retries.start_retry_timer()
return super().send(request, **kwargs)
class TelemetryHttpClient: # TODO: Unify all the http clients in the PySQL Connector
"""Singleton HTTP client for sending telemetry data."""
_instance: Optional["TelemetryHttpClient"] = None
_lock = threading.Lock()
TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT = 3
TELEMETRY_RETRY_DELAY_MIN = 1.0
TELEMETRY_RETRY_DELAY_MAX = 10.0
TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
def __init__(self):
"""Initializes the session and mounts the custom retry adapter."""
retry_policy = DatabricksRetryPolicy(
delay_min=self.TELEMETRY_RETRY_DELAY_MIN,
delay_max=self.TELEMETRY_RETRY_DELAY_MAX,
stop_after_attempts_count=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT,
stop_after_attempts_duration=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION,
delay_default=1.0,
force_dangerous_codes=[],
)
adapter = TelemetryHTTPAdapter(max_retries=retry_policy)
self.session = requests.Session()
self.session.mount("https://", adapter)
self.session.mount("http://", adapter)
@classmethod
def get_instance(cls) -> "TelemetryHttpClient":
"""Get the singleton instance of the TelemetryHttpClient."""
if cls._instance is None:
with cls._lock:
if cls._instance is None:
logger.debug("Initializing singleton TelemetryHttpClient")
cls._instance = TelemetryHttpClient()
return cls._instance
def post(self, url: str, **kwargs) -> requests.Response:
"""
Executes a POST request using the configured session.
This is a blocking call intended to be run in a background thread.
"""
logger.debug("Executing telemetry POST request to: %s", url)
return self.session.post(url, **kwargs)
def close(self):
"""Closes the underlying requests.Session."""
logger.debug("Closing TelemetryHttpClient session.")
self.session.close()
# Clear the instance to allow for re-initialization if needed
with TelemetryHttpClient._lock:
TelemetryHttpClient._instance = None