feat: support env

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -1,3 +1,4 @@
+import os
 import urllib.parse
 import pyarrow as pa
 import importlib.util
@@ -47,6 +48,32 @@ def file_parser_options(**kwargs):
     return kwargs
 
 
+def from_env(**kwargs):
+    """
+    Create an instance of `InfluxDBClient3` using environment variables for configuration.
+
+    This function retrieves the following environment variables:
+      - `INFLUX_HOST`: The hostname or IP address of the InfluxDB server.
+      - `INFLUX_TOKEN`: The authentication token used for accessing the server.
+      - `INFLUX_DATABASE`: The default database for the client operations.
+      - `INFLUX_ORG`: The organization associated with InfluxDB operations.
+
+    If any of these variables are not set, their respective parameters will
+    default to `None`.
+
+    :param kwargs: Additional keyword arguments that will be passed to the
+                   `InfluxDBClient3` constructor for customization.
+    :return: An initialized `InfluxDBClient3` instance.
+    """
+    
+    host = os.getenv("INFLUX_HOST")
+    token = os.getenv("INFLUX_TOKEN")
+    database = os.getenv("INFLUX_DATABASE")
+    org = os.getenv("INFLUX_ORG")
+
+    return InfluxDBClient3(host=host, token=token, database=database, org=org, **kwargs)
+
+
 def _deep_merge(target, source):
     """
     Performs a deep merge of dictionaries or lists,

tests/test_influxdb_client_3.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch
 
-from influxdb_client_3 import InfluxDBClient3
+from influxdb_client_3 import InfluxDBClient3, from_env
 from tests.util import asyncio_run
 from tests.util.mocks import ConstantFlightServer, ConstantData
 
@@ -74,6 +74,40 @@ async def test_query_async(self):
             assert {'data': 'sql_query', 'reference': query, 'value': -1.0} in result_list
             assert {'data': 'query_type', 'reference': 'sql', 'value': -1.0} in result_list
 
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_ORG': 'test_org'})
+    def test_from_env_all_env_vars_set(self):
+        client = from_env()
+        self.assertIsInstance(client, InfluxDBClient3)
+        self.assertEqual(client._client.url, "https://localhost:443")
+        self.assertEqual(client._database, "test_db")
+        self.assertEqual(client._org, "test_org")
+        self.assertEqual(client._token, "test_token")
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_DATABASE': 'test_db'})
+    def test_from_env_partial_env_vars_set(self):
+        client = from_env()
+        self.assertIsInstance(client, InfluxDBClient3)
+        self.assertEqual(client._client.url, "https://localhost:443")
+        self.assertEqual(client._database, "test_db")
+        self.assertEqual(client._org, "default")
+        self.assertIsNone(client._token)
+
+    @patch.dict('os.environ', {}, clear=True)
+    def test_from_env_no_env_vars_set(self):
+        client = from_env()
+        self.assertIsInstance(client, InfluxDBClient3)
+        self.assertIsNotNone(client._client.url)
+        self.assertIsNone(client._database)
+        self.assertIsNone(client._token)
+        self.assertEqual(client._org, "default")
+
+    def test_from_env_with_kargs(self):
+        client = from_env(
+            write_client_options = write_client_options(batch_size=10000),
+        )
+        self.assertIsInstance(client, InfluxDBClient3)
+        self.assertEqual(client._write_client_options['batch_size'], 10000)
 
 if __name__ == '__main__':
     unittest.main()

tests/test_influxdb_client_3_integration.py

@@ -4,12 +4,13 @@
 import string
 import time
 import unittest
+from unittest.mock import patch
 
 import pyarrow
 import pytest
 from pyarrow._flight import FlightError
 
-from influxdb_client_3 import InfluxDBClient3, InfluxDBError, write_client_options, WriteOptions
+from influxdb_client_3 import InfluxDBClient3, InfluxDBError, write_client_options, WriteOptions, from_env
 from tests.util import asyncio_run, lp_to_py_object
 
 
@@ -274,3 +275,20 @@ async def test_verify_query_async(self):
         result_list = result.to_pylist()
         for item in data:
             assert lp_to_py_object(item) in result_list, f"original lp data \"{item}\" should be in result list"
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'https://us-east-1-1.aws.cloud2.influxdata.com',
+                               'INFLUX_TOKEN': 'lDAtMRmhnLp5GjWNVBsieufUb66XZAPxvX3etlmi9wmeq7ispWoL06mwnxmY_BtHKoBhG4lR-c7WfrFgUXy15w==',
+                               'INFLUX_DATABASE': 'bucket0'})
+    def test_from_env(self):
+        c = from_env()
+        with c:
+            id_test = time.time_ns()
+            c.write(f"integration_test_python,type=used value=123.0,id_test={id_test}i")
+
+            sql = 'SELECT * FROM integration_test_python where type=$type and id_test=$id_test'
+            data = c.query(sql, mode="pandas", query_parameters={'type': 'used', 'id_test': id_test})
+
+            self.assertIsNotNone(data)
+            self.assertEqual(1, len(data))
+            self.assertEqual(id_test, data['id_test'][0])
+            self.assertEqual(123.0, data['value'][0])