Add column type to a query result of InfluxDB

redash/query_runner/influx_db.py

@@ -1,6 +1,12 @@
 import logging
 
-from redash.query_runner import BaseQueryRunner, register
+from redash.query_runner import (
+    TYPE_FLOAT,
+    TYPE_INTEGER,
+    TYPE_STRING,
+    BaseQueryRunner,
+    register,
+)
 from redash.utils import json_dumps
 
 logger = logging.getLogger(__name__)
@@ -14,25 +20,36 @@
     enabled = False
 
 
+TYPES_MAP = {
+    str: TYPE_STRING,
+    int: TYPE_INTEGER,
+    float: TYPE_FLOAT,
+}
+
+
+def _get_type(value):
+    return TYPES_MAP.get(type(value), TYPE_STRING)
+
+
 def _transform_result(results):
-    result_columns = []
+    column_names = []
     result_rows = []
 
     for result in results:
         for series in result.raw.get("series", []):
             for column in series["columns"]:
-                if column not in result_columns:
-                    result_columns.append(column)
+                if column not in column_names:
+                    column_names.append(column)
             tags = series.get("tags", {})
             for key in tags.keys():
-                if key not in result_columns:
-                    result_columns.append(key)
+                if key not in column_names:
+                    column_names.append(key)
 
     for result in results:
         for series in result.raw.get("series", []):
             for point in series["values"]:
                 result_row = {}
-                for column in result_columns:
+                for column in column_names:
                     tags = series.get("tags", {})
                     if column in tags:
                         result_row[column] = tags[column]
@@ -42,7 +59,12 @@
                         result_row[column] = value
                 result_rows.append(result_row)
 
-    return json_dumps({"columns": [{"name": c} for c in result_columns], "rows": result_rows})
+    if len(result_rows) > 0:
+        result_columns = [{"name": c, "type": _get_type(result_rows[0][c])} for c in result_rows[0].keys()]
+    else:
+        result_columns = [{"name": c, "type": TYPE_STRING} for c in column_names]
+
+    return json_dumps({"columns": result_columns, "rows": result_rows})
 
 
 class InfluxDB(BaseQueryRunner):

tests/query_runner/test_influx_db.py

@@ -0,0 +1,58 @@
+import json
+
+from influxdb.resultset import ResultSet
+
+from redash.query_runner import (
+    TYPE_FLOAT,
+    TYPE_INTEGER,
+    TYPE_STRING,
+)
+from redash.query_runner.influx_db import _transform_result
+
+raw = {
+    "series": [
+        {
+            "name": "typetest",
+            "columns": ["time", "k1", "v1", "v2"],
+            "values": [
+                ["2023-10-06T13:30:51.323358136Z", "foo", 0.5, 2],
+                ["2023-10-06T13:31:08.882953339Z", "bar", 0.6, 4],
+            ],
+        }
+    ]
+}
+
+raw_no_rows = {"series": [{"name": "typetest", "columns": ["time", "k1", "v1", "v2"], "values": []}]}
+
+
+def test_influxdb_result_types_with_rows():
+    result = ResultSet(raw)
+    transformed = _transform_result([result])
+    expected = {
+        "columns": [
+            {"name": "time", "type": TYPE_STRING},
+            {"name": "k1", "type": TYPE_STRING},
+            {"name": "v1", "type": TYPE_FLOAT},
+            {"name": "v2", "type": TYPE_INTEGER},
+        ],
+        "rows": [
+            {"k1": "foo", "time": "2023-10-06T13:30:51.323358136Z", "v1": 0.5, "v2": 2},
+            {"k1": "bar", "time": "2023-10-06T13:31:08.882953339Z", "v1": 0.6, "v2": 4},
+        ],
+    }
+    assert json.loads(transformed) == expected
+
+
+def test_influxdb_result_types_with_no_rows_are_string():
+    result = ResultSet(raw_no_rows)
+    transformed = _transform_result([result])
+    expected = {
+        "columns": [
+            {"name": "time", "type": TYPE_STRING},
+            {"name": "k1", "type": TYPE_STRING},
+            {"name": "v1", "type": TYPE_STRING},
+            {"name": "v2", "type": TYPE_STRING},
+        ],
+        "rows": [],
+    }
+    assert json.loads(transformed) == expected