feat: support env

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -1,6 +1,7 @@
 import importlib.util
 import os
 import urllib.parse
+from typing import Any
 
 import pyarrow as pa
 
@@ -48,16 +49,19 @@ def file_parser_options(**kwargs):
     """
     return kwargs
 
-
-def from_env(**kwargs):
+# Constants for environment variable names
+INFLUX_HOST = "INFLUX_HOST"
+INFLUX_TOKEN = "INFLUX_TOKEN"
+INFLUX_DATABASE = "INFLUX_DATABASE"
+INFLUX_ORG = "INFLUX_ORG"
+def from_env(**kwargs: Any) -> 'InfluxDBClient3':
     """
     Create an instance of `InfluxDBClient3` using environment variables for configuration.
 
     This function retrieves and validates the following required 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.
-    
     And optional environment variable:
       - `INFLUX_ORG`: The organization associated with InfluxDB operations.
                       Defaults to "default" if not set.
@@ -72,26 +76,26 @@ def from_env(**kwargs):
     :return: An initialized `InfluxDBClient3` instance.
     :raises ValueError: If any required environment variables are not set.
     """
-    
-    invalid_env_vars = []
-    host = os.getenv("INFLUX_HOST")
-    if host is None:
-        invalid_env_vars.append("INFLUX_HOST")
-    
-    token = os.getenv("INFLUX_TOKEN")
-    if token is None:
-        invalid_env_vars.append("INFLUX_TOKEN")
-
-    database = os.getenv("INFLUX_DATABASE")
-    if database is None:
-        invalid_env_vars.append("INFLUX_DATABASE")
-
-    org = os.getenv("INFLUX_ORG")
-
-    if len(invalid_env_vars) > 0:
-        raise ValueError(f"The following environment variables are None or empty: {', '.join(invalid_env_vars)}")
-
-    return InfluxDBClient3(host=host, token=token, database=database, org=org, **kwargs)
+    required_vars = {
+        INFLUX_HOST: os.getenv(INFLUX_HOST),
+        INFLUX_TOKEN: os.getenv(INFLUX_TOKEN),
+        INFLUX_DATABASE: os.getenv(INFLUX_DATABASE)
+    }
+
+    missing_vars = [var for var, value in required_vars.items() if value is None or value == ""]
+    if missing_vars:
+        raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
+
+    org = os.getenv(INFLUX_ORG, "default")
+
+    return InfluxDBClient3(
+        host=required_vars[INFLUX_HOST],
+        token=required_vars[INFLUX_TOKEN],
+        database=required_vars[INFLUX_DATABASE],
+        org=org,
+        **kwargs
+    )
+
 
 
 def _deep_merge(target, source):

tests/test_influxdb_client_3.py

@@ -84,31 +84,19 @@ def test_from_env_all_env_vars_set(self):
         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)
 
+    @patch.dict('os.environ', {'INFLUX_HOST': "", 'INFLUX_TOKEN': "",
+                               'INFLUX_DATABASE': "", 'INFLUX_ORG': ""})
+    def test_from_env_missing_variables(self):
+        with self.assertRaises(ValueError) as context:
+            from_env()
+        self.assertIn("Missing required environment variables", str(context.exception))
 
 if __name__ == '__main__':
     unittest.main()

tests/test_influxdb_client_3_integration.py

@@ -276,9 +276,6 @@ async def test_verify_query_async(self):
         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: