-
Notifications
You must be signed in to change notification settings - Fork 19
refactor: rest client #217
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,10 +1,16 @@ | ||
| import multiprocessing | ||
|
|
||
| import importlib.util | ||
| import json | ||
| import os | ||
| import urllib.parse | ||
| from typing import Any, List, Literal, Optional, TYPE_CHECKING | ||
|
|
||
| import pyarrow as pa | ||
|
|
||
| from influxdb_client_3.version import USER_AGENT | ||
| from influxdb_client_3.write_client._sync import rest_client as rest | ||
|
|
||
| if TYPE_CHECKING: | ||
| import pandas as pd | ||
| import polars as pl | ||
|
|
@@ -14,7 +20,7 @@ | |
| from influxdb_client_3.exceptions import InfluxDBError | ||
| from influxdb_client_3.query.query_api import QueryApi as _QueryApi, QueryApiOptionsBuilder | ||
| from influxdb_client_3.read_file import UploadFile | ||
| from influxdb_client_3.write_client import InfluxDBClient as _InfluxDBClient, WriteOptions, Point | ||
| from influxdb_client_3.write_client import WriteOptions, Point | ||
| from influxdb_client_3.write_client.client.write_api import WriteApi as _WriteApi, SYNCHRONOUS, ASYNCHRONOUS, \ | ||
| PointSettings, DefaultWriteOptions, WriteType | ||
| from influxdb_client_3.write_client.domain.write_precision import WritePrecision | ||
|
|
@@ -189,11 +195,16 @@ def __init__( | |
| org=None, | ||
| database=None, | ||
| token=None, | ||
| auth_scheme=None, | ||
| enable_gzip=False, | ||
| gzip_threshold=None, | ||
| write_client_options=None, | ||
| flight_client_options=None, | ||
| write_port_overwrite=None, | ||
| query_port_overwrite=None, | ||
| disable_grpc_compression=False, | ||
| point_settings=None, | ||
| debug=False, | ||
| **kwargs): | ||
| """ | ||
| Initialize an InfluxDB client. | ||
|
|
@@ -212,6 +223,14 @@ def __init__( | |
| :type flight_client_options: dict[str, any] | ||
| :param disable_grpc_compression: Disable gRPC compression for Flight query responses. Default is False. | ||
| :type disable_grpc_compression: bool | ||
| :param point_settings The settings for Points | ||
| :type point_settings: PointSettings | ||
| :param debug: enable verbose logging of http requests | ||
| :type debug: bool | ||
| :param enable_gzip: Enable GZIP compression for write requests. | ||
| :type enable_gzip: bool | ||
| :param gzip_threshold: Minimum payload size (bytes) to trigger GZIP when enable_gzip is True. | ||
| :type gzip_threshold: int | ||
| :key auth_scheme: token authentication scheme. Set to "Bearer" for Edge. | ||
| :key bool verify_ssl: Set this to false to skip verifying SSL certificate when calling API from https server. | ||
| :key str ssl_ca_cert: Set this to customize the certificate file to verify the peer. | ||
|
|
@@ -235,6 +254,10 @@ def __init__( | |
| :key bool write_no_sync: disable sync confirmation on V3 API endpoint writes. | ||
| :key list[str] profilers: list of enabled Flux profilers | ||
| """ | ||
| for key in ["host", "token", "database"]: | ||
| if locals().get(key) is None: | ||
| raise Exception(f"The '{key}' key is required") | ||
|
|
||
| self._org = org if org is not None else "default" | ||
| self._database = database | ||
| self._token = token | ||
|
|
@@ -293,14 +316,49 @@ def __init__( | |
| if write_port_overwrite is not None: | ||
| port = write_port_overwrite | ||
|
|
||
| self._client = _InfluxDBClient( | ||
| url=f"{scheme}://{hostname}:{port}", | ||
| token=self._token, | ||
| auth_schema = 'Token' if auth_scheme is None else auth_scheme | ||
| default_header = { | ||
| 'User-Agent': USER_AGENT | ||
| } | ||
| if self._token is not None: | ||
|
NguyenHoangSon96 marked this conversation as resolved.
|
||
| default_header['Authorization'] = f'{auth_schema} {self._token}' | ||
| self.base_url = f"{scheme}://{hostname}:{port}" | ||
| self.default_header = default_header | ||
|
NguyenHoangSon96 marked this conversation as resolved.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I see that the default value for I might recommend a guard for arguments that should always be defined at the start of the for key in ["host", "token", "database"]:
if locals().get(key) is None:
raise Exception(f"The '{key}' key is required")This is of course unless I've overlooked another type enforcement mechanism in Python, of which I'm not yet aware. Currently this is not an issue for unit tests using a mock, since the mock essentially just makes use of or ignores the 'None' part of the URL. However, outside of a testing framework, this should not be allowed.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| self.rest_client = rest.RestClient( | ||
| base_url=self.base_url, | ||
| default_header=default_header, | ||
|
NguyenHoangSon96 marked this conversation as resolved.
|
||
| verify_ssl=kwargs.get('verify_ssl', True), | ||
| ssl_ca_cert=kwargs.get('ssl_ca_cert', None), | ||
| cert_file=kwargs.get('cert_file', None), | ||
| cert_key_file=kwargs.get('cert_key_file', None), | ||
| cert_key_password=kwargs.get('cert_key_password', None), | ||
| ssl_context=kwargs.get('ssl_context', None), | ||
| proxy=kwargs.get('proxy', None), | ||
| proxy_headers=kwargs.get('proxy_headers', None), | ||
| retries=kwargs.get('retries', False), | ||
| debug=debug, | ||
| connection_pool_maxsize=kwargs.get('connection_pool_maxsize', multiprocessing.cpu_count() * 5,) | ||
| ) | ||
|
|
||
| if point_settings is None: | ||
| point_settings = PointSettings() | ||
|
|
||
| # Keep WriteOptions.timeout in sync with the resolved write_timeout | ||
| if isinstance(self._write_client_options, dict) and self._write_client_options.get("write_options") is not None: | ||
| self._write_client_options["write_options"].timeout = write_timeout | ||
|
|
||
| self._write_api = _WriteApi( | ||
|
NguyenHoangSon96 marked this conversation as resolved.
Copilot marked this conversation as resolved.
|
||
| bucket=self._database, | ||
| org=self._org, | ||
| gzip_threshold=gzip_threshold, | ||
| enable_gzip=enable_gzip, | ||
| auth_scheme=auth_scheme, | ||
| timeout=write_timeout, | ||
| **kwargs) | ||
|
|
||
| self._write_api = _WriteApi(influxdb_client=self._client, **self._write_client_options) | ||
| default_header=default_header, | ||
| rest_client=self.rest_client, | ||
| point_settings=point_settings, | ||
| **self._write_client_options | ||
| ) | ||
|
|
||
| if query_port_overwrite is not None: | ||
| port = query_port_overwrite | ||
|
|
@@ -658,32 +716,32 @@ async def query_async(self, query: str, language: str = "sql", mode: str = "all" | |
| except ArrowException as e: | ||
| raise InfluxDB3ClientQueryError(f"Error while executing query: {e}") | ||
|
|
||
| def get_server_version(self) -> str: | ||
| def get_server_version(self) -> Optional[str]: | ||
| """ | ||
| Get the version of the connected InfluxDB server. | ||
| Retrieves the server version by querying the designated endpoint and | ||
| extracting the version information from either response headers or | ||
| the response body. | ||
|
|
||
| This method makes a ping request to the server and extracts the version information | ||
| from either the response headers or response body. | ||
| This method interacts with a REST API endpoint to fetch the server's | ||
| version details, which might be stored in a specific HTTP header or | ||
| available in the response body as part of a JSON structure. | ||
|
|
||
| :return: The version string of the InfluxDB server. | ||
| :rtype: str | ||
| :return: The version string of the server if available, otherwise None. | ||
| :rtype: Optional[str] | ||
| """ | ||
| version = None | ||
| (resp_body, _, header) = self._client.api_client.call_api( | ||
| resource_path="/ping", | ||
| method="GET", | ||
| response_type=object | ||
| ) | ||
|
|
||
| for key, value in header.items(): | ||
| resp = self.rest_client.request(url='/ping', method="GET", headers=self.default_header) | ||
| for key, value in resp.getheaders().items(): | ||
| if key.lower() == "x-influxdb-version": | ||
| version = value | ||
| break | ||
|
|
||
| if version is None and isinstance(resp_body, dict): | ||
| version = resp_body['version'] | ||
| return value | ||
|
|
||
| return version | ||
| string_body = resp.get_string_body() | ||
| if not string_body: | ||
| return None | ||
| try: | ||
| data = json.loads(string_body) | ||
| except (ValueError, TypeError): | ||
| return None | ||
| return data.get('version') | ||
|
|
||
| def flush(self): | ||
| """ | ||
|
|
@@ -700,9 +758,12 @@ def flush(self): | |
|
|
||
| def close(self): | ||
| """Close the client and clean up resources.""" | ||
| self._write_api.close() | ||
| self._query_api.close() | ||
| self._client.close() | ||
| if self._write_api is not None: | ||
| self._write_api.close() | ||
| if self._query_api is not None: | ||
|
NguyenHoangSon96 marked this conversation as resolved.
|
||
| self._query_api.close() | ||
| if self.rest_client is not None: | ||
| self.rest_client.close() | ||
|
|
||
| def __enter__(self): | ||
| return self | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.