feat: add supports for array expressions in query parameters (influxdata#319)

bednar · web-flow · commit e2d791109868 · 2021-08-23T09:37:12.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,8 @@
 ## 1.21.0 [unreleased]
 
+### Features
+1. [#xxx](https://github.com/influxdata/influxdb-client-python/pull/xxx): Add supports for array expressions in query parameters
+ 
 ## 1.20.0 [2021-08-20]
 
 ### Features
diff --git a/influxdb_client/client/query_api.py b/influxdb_client/client/query_api.py
@@ -7,10 +7,10 @@
 import codecs
 import csv
 from datetime import datetime, timedelta
-from typing import List, Generator, Any
+from typing import List, Generator, Any, Union, Iterable
 
 from influxdb_client import Dialect, IntegerLiteral, BooleanLiteral, FloatLiteral, DateTimeLiteral, StringLiteral, \
-    VariableAssignment, Identifier, OptionStatement, File, DurationLiteral, Duration, UnaryExpression, \
+    VariableAssignment, Identifier, OptionStatement, File, DurationLiteral, Duration, UnaryExpression, Expression, \
     ImportDeclaration, MemberAssignment, MemberExpression, ArrayExpression
 from influxdb_client import Query, QueryService
 from influxdb_client.client.flux_csv_parser import FluxCsvParser, FluxSerializationMode
@@ -203,35 +203,43 @@ def _params_to_extern_ast(params: dict) -> List['OptionStatement']:
 
         statements = []
         for key, value in params.items():
-            if value is None:
+            expression = QueryApi._parm_to_extern_ast(value)
+            if expression is None:
                 continue
 
-            if isinstance(value, bool):
-                literal = BooleanLiteral("BooleanLiteral", value)
-            elif isinstance(value, int):
-                literal = IntegerLiteral("IntegerLiteral", str(value))
-            elif isinstance(value, float):
-                literal = FloatLiteral("FloatLiteral", value)
-            elif isinstance(value, datetime):
-                value = get_date_helper().to_utc(value)
-                literal = DateTimeLiteral("DateTimeLiteral", value.strftime('%Y-%m-%dT%H:%M:%S.%fZ'))
-            elif isinstance(value, timedelta):
-                _micro_delta = int(value / timedelta(microseconds=1))
-                if _micro_delta < 0:
-                    literal = UnaryExpression("UnaryExpression", argument=DurationLiteral("DurationLiteral", [
-                        Duration(magnitude=-_micro_delta, unit="us")]), operator="-")
-                else:
-                    literal = DurationLiteral("DurationLiteral", [Duration(magnitude=_micro_delta, unit="us")])
-            elif isinstance(value, str):
-                literal = StringLiteral("StringLiteral", str(value))
-            else:
-                literal = value
-
             statements.append(OptionStatement("OptionStatement",
                                               VariableAssignment("VariableAssignment", Identifier("Identifier", key),
-                                                                 literal)))
+                                                                 expression)))
         return statements
 
+    @staticmethod
+    def _parm_to_extern_ast(value) -> Union[Expression, None]:
+        if value is None:
+            return None
+        if isinstance(value, bool):
+            return BooleanLiteral("BooleanLiteral", value)
+        elif isinstance(value, int):
+            return IntegerLiteral("IntegerLiteral", str(value))
+        elif isinstance(value, float):
+            return FloatLiteral("FloatLiteral", value)
+        elif isinstance(value, datetime):
+            value = get_date_helper().to_utc(value)
+            return DateTimeLiteral("DateTimeLiteral", value.strftime('%Y-%m-%dT%H:%M:%S.%fZ'))
+        elif isinstance(value, timedelta):
+            _micro_delta = int(value / timedelta(microseconds=1))
+            if _micro_delta < 0:
+                return UnaryExpression("UnaryExpression", argument=DurationLiteral("DurationLiteral", [
+                    Duration(magnitude=-_micro_delta, unit="us")]), operator="-")
+            else:
+                return DurationLiteral("DurationLiteral", [Duration(magnitude=_micro_delta, unit="us")])
+        elif isinstance(value, str):
+            return StringLiteral("StringLiteral", str(value))
+        elif isinstance(value, Iterable):
+            return ArrayExpression("ArrayExpression",
+                                   elements=list(map(lambda it: QueryApi._parm_to_extern_ast(it), value)))
+        else:
+            return value
+
     @staticmethod
     def _build_flux_ast(params: dict = None, profilers: List[str] = None):
 
diff --git a/tests/test_QueryApi.py b/tests/test_QueryApi.py
@@ -5,7 +5,7 @@
 from dateutil.tz import tzutc
 from httpretty import httpretty
 
-from influxdb_client import QueryApi, DurationLiteral, Duration, CallExpression, Expression, UnaryExpression, \
+from influxdb_client import QueryApi, DurationLiteral, Duration, CallExpression, UnaryExpression, \
     Identifier, InfluxDBClient
 from influxdb_client.client.query_api import QueryOptions
 from influxdb_client.client.util.date_utils import get_date_helper
@@ -47,7 +47,7 @@ def test_query_flux_table(self):
         val_count = 0
         for table in tables:
             for row in table:
-                for cell in row.values:
+                for _ in row.values:
                     val_count += 1
 
         print("Values count: ", val_count)
@@ -61,7 +61,7 @@ def test_query_flux_csv(self):
 
         val_count = 0
         for row in csv_result:
-            for cell in row:
+            for _ in row:
                 val_count += 1
 
         print("Values count: ", val_count)
@@ -98,7 +98,7 @@ def test_query_ast(self):
 
         val_count = 0
         for row in csv_result:
-            for cell in row:
+            for _ in row:
                 val_count += 1
 
         print("Values count: ", val_count)
@@ -263,7 +263,39 @@ def test_parameter_ast(self):
                      }
                  ],
                  "imports": []
-             }]]
+             }], ["arrayParam", ["bar1", "bar2", "bar3"],
+                  {
+                      "body": [
+                          {
+                              "assignment": {
+                                  "id": {
+                                      "name": "arrayParam",
+                                      "type": "Identifier"
+                                  },
+                                  "init": {
+                                      "elements": [
+                                          {
+                                              "type": "StringLiteral",
+                                              "value": "bar1"
+                                          },
+                                          {
+                                              "type": "StringLiteral",
+                                              "value": "bar2"
+                                          },
+                                          {
+                                              "type": "StringLiteral",
+                                              "value": "bar3"
+                                          }
+                                      ],
+                                      "type": "ArrayExpression"
+                                  },
+                                  "type": "VariableAssignment"
+                              },
+                              "type": "OptionStatement"
+                          }
+                      ],
+                      "imports": []
+                  }]]
 
         for data in test_data:
             param = {data[0]: data[1]}
@@ -290,7 +322,7 @@ def test_query_profiler_enabled(self):
         for table in csv_result:
             self.assertFalse(any(filter(lambda column: (column.default_value == "_profiler"), table.columns)))
             for flux_record in table:
-                self.assertFalse( flux_record["_measurement"].startswith("profiler/"))
+                self.assertFalse(flux_record["_measurement"].startswith("profiler/"))
 
         records = self.client.query_api().query_stream(query=q, params=p)
 
