dsquery: Format code and satisfy linter

amotl · amotl · commit 6851cdd252a1 · 2023-09-15T10:55:46.000-07:00
diff --git a/examples/datasource-smartquery.py b/examples/datasource-smartquery.py
@@ -42,7 +42,7 @@
         --query "sum by (status) (count_over_time({filename=~\".+/json_access.+\", job=~\".*\", instance=~\".*\"} | json |  __error__=\"\" [5m]))" \
         --attr "legendFormat:HTTP Status {{status}}"
 
-"""
+"""  # noqa: E501
 import json
 import logging
 from optparse import OptionParser
@@ -69,9 +69,8 @@ def run(grafana: GrafanaApi):
     if not options.query:
         parser.error("Option --query required")
 
-    attributes = {
-    }
-    if options.attrs is not None and len(options.attrs) > 0 :
+    attributes = {}
+    if options.attrs is not None and len(options.attrs) > 0:
         for attr in options.attrs:
             (key, val) = attr.split(":")
             attributes[key] = val
diff --git a/grafana_client/elements/datasource.py b/grafana_client/elements/datasource.py
@@ -308,8 +308,13 @@ def series(self, datasource_id, match, start, end, access="proxy"):
         r = self.client.POST(post_series_path, data={"match[]": match, "start": start, "end": end})
         return r
 
-#**********************************************************************************
-    def smartquery(self, datasource: Union[DatasourceIdentifier, Dict], expression: str, attrs: Optional[dict] = None, request: Optional[dict] = None):
+    def smartquery(
+        self,
+        datasource: Union[DatasourceIdentifier, Dict],
+        expression: str,
+        attrs: Optional[dict] = None,
+        request: Optional[dict] = None,
+    ):
         """
         Send a query to the designated data source and return its response.
 
@@ -333,15 +338,15 @@ def smartquery(self, datasource: Union[DatasourceIdentifier, Dict], expression:
             model = {
                 "refId": "test",
             }
-            if expression is not None and ( attrs is None or (attrs is not None and 'query' not in attrs)):
-                model['query'] = expression
+            if expression is not None and (attrs is None or (attrs is not None and "query" not in attrs)):
+                model["query"] = expression
             if attrs is not None:
                 model.update(attrs)
             request = query_factory(datasource, model)
 
         # Compute request method, body, and endpoint.
         if "method" in request and isinstance(request["method"], str):
-            if request['method'] == 'POST':
+            if request["method"] == "POST":
                 send_request = self.client.POST
             else:
                 send_request = self.client.GET
@@ -351,16 +356,16 @@ def smartquery(self, datasource: Union[DatasourceIdentifier, Dict], expression:
         # Certain data sources like InfluxDB 1.x, still use the `/datasources/proxy` route.
         if datasource_type == "influxdb" and datasource_dialect == "InfluxQL":
             url = f"/datasources/proxy/{datasource_id}/query"
-            url += '?' + urlencode(request['params'])
+            url += "?" + urlencode(request["params"])
             request_kwargs = {"data": request["data"]}
 
         elif datasource_type == "graphite":
             url = f"/datasources/proxy/{datasource_id}/render"
             request_kwargs = {"data": request["data"]}
 
         # This case is very special. It is used for Elasticsearch and Testdata.
-        elif 'url' in request and model['url'].startswith("url://"):
-            url = model['url'].replace("url://", "")
+        elif "url" in request and model["url"].startswith("url://"):
+            url = model["url"].replace("url://", "")
             url = url.format(
                 datasource_id=datasource.get("id"),
                 datasource_uid=datasource.get("uid"),
@@ -370,8 +375,12 @@ def smartquery(self, datasource: Union[DatasourceIdentifier, Dict], expression:
             send_request = self.client.GET
 
         elif datasource_type in ("prometheus", "loki") and LooseVersion(self.api.version) <= VERSION_7:
-            if "queries" in request["data"] and len(request["data"]["queries"]) > 0 \
-                and "instant" in request["data"]["queries"][0] and request["data"]["queries"][0]["instant"]:
+            if (
+                "queries" in request["data"]
+                and len(request["data"]["queries"]) > 0
+                and "instant" in request["data"]["queries"][0]
+                and request["data"]["queries"][0]["instant"]
+            ):
                 return self.query(
                     datasource.get("id"),
                     request["expr"],
diff --git a/grafana_client/knowledge.py b/grafana_client/knowledge.py
@@ -4,8 +4,8 @@
 Knowledgebase module for filling the missing gaps to access the Grafana API
 more efficiently.
 """
-from typing import Dict, Optional, Union
 from datetime import datetime
+from typing import Dict, Optional, Union
 
 from grafana_client.model import DatasourceModel
 
@@ -143,23 +143,16 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
         query = expression
     elif datasource_type == "grafana-simple-json-datasource":
         query = expression
-    #************************************************************************
+
     elif datasource_type == "graphite":
-        query = {
-            "target": expression,
-            "from": "-5m",
-            "until": "now",
-            "format": "json",
-            "maxDataPoints": 300
-        }
-        if 'time_from' in model:
-            query['from']=model['time_from']
-        if 'time_to' in model:
-            query['until']=model['time_to']
+        query = {"target": expression, "from": "-5m", "until": "now", "format": "json", "maxDataPoints": 300}
+        if "time_from" in model:
+            query["from"] = model["time_from"]
+        if "time_to" in model:
+            query["until"] = model["time_to"]
 
         request["data"] = query
 
-    #************************************************************************
     elif datasource_type == "influxdb":
         dialect = datasource["jsonData"].get("version", "InfluxQL")
         query = {
@@ -177,13 +170,12 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
                 }
             )
             # this drive the how timestamp are rendered in result (string by default, or in milliseconds ms)
-            request["params"] =  { "epoch": "ms" }
-            if 'database' in datasource:
-                request["params"].update({ "db": datasource['database'] })
+            request["params"] = {"epoch": "ms"}
+            if "database" in datasource:
+                request["params"].update({"db": datasource["database"]})
             request["data"] = query
 
         elif dialect == "Flux":
-
             query = {
                 "datasource": {
                     "type": datasource["type"],
@@ -195,17 +187,25 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
             }
 
             attrs = [
-                { "name": "intervalMs", "default": 30000, },
-                { "name": "maxDataPoints", "default": 1441, },
-                { "name": "refId", "default": "test", },
+                {
+                    "name": "intervalMs",
+                    "default": 30000,
+                },
+                {
+                    "name": "maxDataPoints",
+                    "default": 1441,
+                },
+                {
+                    "name": "refId",
+                    "default": "test",
+                },
             ]
         else:
             raise KeyError(f"InfluxDB dialect '{dialect}' unknown")
 
     elif datasource_type == "jaeger":
         query = {}
 
-    #************************************************************************
     elif datasource_type == "loki":
         query = {
             "datasource": {
@@ -218,20 +218,40 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
         }
 
         attrs = [
-            { "name": "intervalMs", "default": 60000, },
-            { "name": "legendFormat", "default": "", },
-            { "name": "maxLines", "default": 1000, },
-            { "name": "maxDataPoints", "default": 1441, },
-            { "name": "queryType", "default": "range", },
-            { "name": "refId", "default": "test", },
-            { "name": "resolution", "default": 1, },
+            {
+                "name": "intervalMs",
+                "default": 60000,
+            },
+            {
+                "name": "legendFormat",
+                "default": "",
+            },
+            {
+                "name": "maxLines",
+                "default": 1000,
+            },
+            {
+                "name": "maxDataPoints",
+                "default": 1441,
+            },
+            {
+                "name": "queryType",
+                "default": "range",
+            },
+            {
+                "name": "refId",
+                "default": "test",
+            },
+            {
+                "name": "resolution",
+                "default": 1,
+            },
         ]
-    
+
     elif datasource_type == "opentsdb":
         query = {}
 
-    #************************************************************************
-    elif datasource_type in ( "postgres", "mssql", "mysql" ):
+    elif datasource_type in ("postgres", "mssql", "mysql"):
         query = {
             "datasource": {
                 "type": datasource["type"],
@@ -241,16 +261,26 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
             "rawSql": expression,
         }
         attrs = [
-            { "name": "format", "default": "time_series",
-                "choices": [ "time_series", "table"],
+            {
+                "name": "format",
+                "default": "time_series",
+                "choices": ["time_series", "table"],
                 # "version": "8.0.0"
             },
-            { "name": "intervalMs", "default": 30000, },
-            { "name": "maxDataPoints", "default": 1441, },
-            { "name": "refId", "default": None, },
+            {
+                "name": "intervalMs",
+                "default": 30000,
+            },
+            {
+                "name": "maxDataPoints",
+                "default": 1441,
+            },
+            {
+                "name": "refId",
+                "default": None,
+            },
         ]
 
-    #************************************************************************
     elif datasource_type == "prometheus":
         query = {
             "datasource": {
@@ -263,21 +293,56 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
         }
 
         attrs = [
-            { "name": "format", "default": "time_series",
-                "choices": [ "time_series", "table", "heatmap"],
+            {
+                "name": "format",
+                "default": "time_series",
+                "choices": ["time_series", "table", "heatmap"],
                 # "version": "8.0.0"
             },
-            { "name": "instant", "default": False, },
-            { "name": "interval", "default": "", },
-            { "name": "intervalFactor", "default": None, },
-            { "name": "intervalMs", "default": 30000, },
-            { "name": "legendFormat", "default": "", },
-            { "name": "maxDataPoints", "default": None, },
-            { "name": "queryType", "default": "timeSeriesQuery", },
-            { "name": "refId", "default": "test", },
-            { "name": "requestId", "default": "0test", },
-            { "name": "step", "default": 300, },
-            { "name": "utcOffsetSec", "default": 0, },
+            {
+                "name": "instant",
+                "default": False,
+            },
+            {
+                "name": "interval",
+                "default": "",
+            },
+            {
+                "name": "intervalFactor",
+                "default": None,
+            },
+            {
+                "name": "intervalMs",
+                "default": 30000,
+            },
+            {
+                "name": "legendFormat",
+                "default": "",
+            },
+            {
+                "name": "maxDataPoints",
+                "default": None,
+            },
+            {
+                "name": "queryType",
+                "default": "timeSeriesQuery",
+            },
+            {
+                "name": "refId",
+                "default": "test",
+            },
+            {
+                "name": "requestId",
+                "default": "0test",
+            },
+            {
+                "name": "step",
+                "default": 300,
+            },
+            {
+                "name": "utcOffsetSec",
+                "default": 0,
+            },
         ]
 
     elif datasource_type == "simpod-json-datasource":
@@ -304,20 +369,20 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
             if value is not None:
                 query[attr["name"]] = value
 
-        if 'time_from' not in model or 'time_to' not in model:
+        if "time_from" not in model or "time_to" not in model:
             now = datetime.now()
-            if 'time_from' not in model:
-                model['time_from'] = int( now.timestamp() ) - 5 *60
-            if 'time_to' not in model:
-                model['time_to'] = int( now.timestamp() )
+            if "time_from" not in model:
+                model["time_from"] = int(now.timestamp()) - 5 * 60
+            if "time_to" not in model:
+                model["time_to"] = int(now.timestamp())
 
-        if 'instant' in query and query['instant']:
-            model['time_from'] = model['time_to']
+        if "instant" in query and query["instant"]:
+            model["time_from"] = model["time_to"]
 
         payload = {
             "queries": [query],
-            "from": str(model['time_from']*1000),
-            "to": str(model['time_to']*1000),
+            "from": str(model["time_from"] * 1000),
+            "to": str(model["time_to"] * 1000),
         }
         request["data"] = payload
 
@@ -330,12 +395,13 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
     "elasticsearch": "url:///datasources/proxy/{datasource_id}/{database_name}/_mapping",
     "fetzerch-sunandmoon-datasource": "url:///datasources/uid/{datasource_uid}",
     "grafana-simple-json-datasource": "url:///datasources/proxy/{datasource_id}",
-    "graphite": "random-walk.count;dc=asia-1;app=collector;server=000",  # from play.grafana.org explore Grafite datasource.
+    # From play.grafana.org, to explore Graphite datasource.
+    "graphite": "random-walk.count;dc=asia-1;app=collector;server=000",
     "influxdb": "SHOW RETENTION POLICIES on _internal",
     "influxdb+influxql": "SHOW RETENTION POLICIES on _internal",
     "influxdb+flux": "buckets()",
     "jaeger": "url:///datasources/proxy/{datasource_id}/api/services",
-    "loki": "count_over_time({job=~\".+\"} [5m])",
+    "loki": 'count_over_time({job=~".+"} [5m])',
     "mssql": "SELECT 1",
     "mysql": "SELECT 1",
     "opentsdb": "url:///datasources/proxy/{datasource_id}/api/suggest?type=metrics&q=cpu&max=100",
diff --git a/test/elements/test_datasource_base.py b/test/elements/test_datasource_base.py
