Skip to content

Commit d904983

Browse files
committed
add proxies
Signed-off-by: Eric <eric@epsilla.com>
1 parent f7f3cef commit d904983

File tree

1 file changed

+83
-15
lines changed

1 file changed

+83
-15
lines changed

pyepsilla/cloud/client.py

Lines changed: 83 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515

1616

1717
class Client(object):
18-
def __init__(self, project_id: str, api_key: str, headers: dict = None):
18+
def __init__(
19+
self, project_id: str, api_key: str, headers: dict = None, proxies: dict = None
20+
):
1921
self._project_id = project_id
2022
self._apikey = api_key
2123
self._baseurl = f"https://dispatch.epsilla.com/api/v3/project/{self._project_id}" # type: ignore
@@ -25,16 +27,21 @@ def __init__(self, project_id: str, api_key: str, headers: dict = None):
2527
"Connection": "close",
2628
"X-API-Key": api_key,
2729
}
30+
self._proxy = None
31+
if proxies is not None:
32+
self._proxy = proxies
2833
if headers is not None:
2934
self._header.update(headers)
3035
self._db_id = None
3136

3237
def validate(self):
38+
req_url = f"{self._baseurl}/vectordb/list"
3339
resp = requests.get(
34-
url=f"{self._baseurl}/vectordb/list",
40+
url=req_url,
3541
data=None,
3642
headers=self._header,
3743
verify=False,
44+
proxies=self._proxy,
3845
)
3946
data = resp.json()
4047
resp.close()
@@ -44,7 +51,13 @@ def validate(self):
4451
def get_db_list(self):
4552
db_list = []
4653
req_url = f"{self._baseurl}/vectordb/list"
47-
resp = requests.get(url=req_url, data=None, headers=self._header, verify=False)
54+
resp = requests.get(
55+
url=req_url,
56+
data=None,
57+
headers=self._header,
58+
verify=False,
59+
proxies=self._proxy,
60+
)
4861
status_code = resp.status_code
4962
body = resp.json()
5063
if status_code == 200 and body["statusCode"] == 200:
@@ -56,7 +69,13 @@ def get_db_list(self):
5669
def load_db(self, db_name: str, db_path: str):
5770
db_id = db_name.lstrip("db_").replace("_", "-")
5871
req_url = f"{self._baseurl}/vectordb/{db_id}/load"
59-
resp = requests.post(url=req_url, data=None, headers=self._header, verify=False)
72+
resp = requests.post(
73+
url=req_url,
74+
data=None,
75+
headers=self._header,
76+
verify=False,
77+
proxies=self._proxy,
78+
)
6079
status_code = resp.status_code
6180
body = resp.json()
6281
resp.close()
@@ -69,7 +88,13 @@ def use_db(self, db_name: str):
6988

7089
def get_db_info(self, db_id: str):
7190
req_url = f"{self._baseurl}/vectordb/{db_id}"
72-
resp = requests.get(url=req_url, data=None, headers=self._header, verify=False)
91+
resp = requests.get(
92+
url=req_url,
93+
data=None,
94+
headers=self._header,
95+
verify=False,
96+
proxies=self._proxy,
97+
)
7398
status_code = resp.status_code
7499
body = resp.json()
75100
resp.close()
@@ -80,7 +105,11 @@ def get_db_statistics(self, db_id: str):
80105
req_url = f"{self._baseurl}/vectordb/{db_id}/statistics"
81106
req_data = None
82107
resp = requests.get(
83-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
108+
url=req_url,
109+
data=json.dumps(req_data),
110+
headers=self._header,
111+
verify=False,
112+
proxies=self._proxy,
84113
)
85114
status_code = resp.status_code
86115
body = resp.json()
@@ -106,7 +135,12 @@ def vectordb(self, db_id: str):
106135
status_code, resp = self.get_db_info(db_id=db_id)
107136
if resp["statusCode"] == 200:
108137
return Vectordb(
109-
self._project_id, db_id, self._apikey, resp["result"]["public_endpoint"]
138+
self._project_id,
139+
db_id,
140+
self._apikey,
141+
resp["result"]["public_endpoint"],
142+
self._header,
143+
self._proxy,
110144
)
111145
else:
112146
print(resp)
@@ -122,13 +156,17 @@ def __init__(
122156
api_key: str,
123157
public_endpoint: str,
124158
headers: dict = None,
159+
proxies: dict = None,
125160
):
126161
self._project_id = project_id
127162
self._db_id = db_id
128163
self._api_key = api_key
129164
self._public_endpoint = public_endpoint
130165
self._baseurl = f"https://{self._public_endpoint}/api/v3/project/{self._project_id}/vectordb/{self._db_id}"
131166
self._header = {"Content-type": "application/json", "X-API-Key": self._api_key}
167+
self._proxy = None
168+
if proxies is not None:
169+
self._proxy = proxies
132170
if headers is not None:
133171
self._header.update(headers)
134172

@@ -137,7 +175,9 @@ def list_tables(self):
137175
if self._db_id is None:
138176
raise Exception("[ERROR] db_id is None!")
139177
req_url = f"{self._baseurl}/table/list"
140-
resp = requests.get(url=req_url, headers=self._header, verify=False)
178+
resp = requests.get(
179+
url=req_url, headers=self._header, verify=False, proxies=self._proxy
180+
)
141181
status_code = resp.status_code
142182
body = resp.json()
143183
resp.close()
@@ -160,7 +200,11 @@ def create_table(
160200
if indices is not None:
161201
req_data["indices"] = indices
162202
resp = requests.post(
163-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
203+
url=req_url,
204+
data=json.dumps(req_data),
205+
headers=self._header,
206+
verify=False,
207+
proxies=self._proxy,
164208
)
165209
status_code = resp.status_code
166210
body = resp.json()
@@ -175,7 +219,11 @@ def drop_table(self, table_name: str):
175219
req_url = f"{self._baseurl}/table/delete?table_name={table_name}"
176220
req_data = {}
177221
resp = requests.delete(
178-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
222+
url=req_url,
223+
data=json.dumps(req_data),
224+
headers=self._header,
225+
verify=False,
226+
proxies=self._proxy,
179227
)
180228
status_code = resp.status_code
181229
body = resp.json()
@@ -188,7 +236,11 @@ def insert(self, table_name: str, records: list[dict]):
188236
req_url = f"{self._baseurl}/data/insert"
189237
req_data = {"table": table_name, "data": records}
190238
resp = requests.post(
191-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
239+
url=req_url,
240+
data=json.dumps(req_data),
241+
headers=self._header,
242+
verify=False,
243+
proxies=self._proxy,
192244
)
193245
status_code = resp.status_code
194246
body = resp.json()
@@ -200,7 +252,11 @@ def upsert(self, table_name: str, records: list[dict]):
200252
req_url = f"{self._baseurl}/data/insert"
201253
req_data = {"table": table_name, "data": records, "upsert": True}
202254
resp = requests.post(
203-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
255+
url=req_url,
256+
data=json.dumps(req_data),
257+
headers=self._header,
258+
verify=False,
259+
proxies=self._proxy,
204260
)
205261
status_code = resp.status_code
206262
body = resp.json()
@@ -254,7 +310,11 @@ def query(
254310
req_data["facets"] = facets
255311

256312
resp = requests.post(
257-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
313+
url=req_url,
314+
data=json.dumps(req_data),
315+
headers=self._header,
316+
verify=False,
317+
proxies=self._proxy,
258318
)
259319
status_code = resp.status_code
260320
body = resp.json()
@@ -295,7 +355,11 @@ def delete(
295355
req_data["filter"] = filter
296356

297357
resp = requests.post(
298-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
358+
url=req_url,
359+
data=json.dumps(req_data),
360+
headers=self._header,
361+
verify=False,
362+
proxies=self._proxy,
299363
)
300364
status_code = resp.status_code
301365
body = resp.json()
@@ -352,7 +416,11 @@ def get(
352416

353417
req_url = f"{self._baseurl}/data/get"
354418
resp = requests.post(
355-
url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
419+
url=req_url,
420+
data=json.dumps(req_data),
421+
headers=self._header,
422+
verify=False,
423+
proxies=self._proxy,
356424
)
357425
status_code = resp.status_code
358426
body = resp.json()

0 commit comments

Comments
 (0)