dsquery: new smartquery function in conjonction with query_factory

Author: peekjef72

grafana_client/elements/datasource.py

@@ -4,6 +4,7 @@
 import warnings
 from distutils.version import LooseVersion
 from typing import Dict, Optional, Tuple, Union
+from urllib.parse import urlencode
 
 from requests import ReadTimeout
 
@@ -306,7 +307,7 @@ 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, store: Optional[str] = None):
+    def smartquery_old(self, datasource: Union[DatasourceIdentifier, Dict], expression: str, store: Optional[str] = None):
         """
         Send a query to the designated data source and return its response.
 
@@ -378,6 +379,85 @@ def smartquery(self, datasource: Union[DatasourceIdentifier, Dict], expression:
                 f"Reason: {ex}. Response: {ex.response or '<empty>'}"
             )
             raise
+#**********************************************************************************
+    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.
+
+        TODO: This is by far not complete. The `query_factory` function has to
+            be made more elaborate in order to query different data source
+                types.
+        """
+
+        if isinstance(datasource, DatasourceIdentifier):
+            datasource = self.get(datasource)
+
+        datasource_id = datasource["id"]
+        datasource_type = datasource["type"]
+        datasource_dialect = datasource.get("jsonData", {}).get("version", "InfluxQL")
+        access_type = datasource["access"]
+
+        # Sanity checks.
+        if not request and not expression:
+            raise ValueError("request or expression must be given")
+        elif not request:
+            model = {
+                "refId": "test",
+            }
+            if expression is not None and 'query' not in attrs:
+                model['query'] = expression
+            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':
+                send_request = self.client.POST
+            else:
+                send_request = self.client.GET
+
+        logger.info(f"Submitting request: {request}")
+
+        # 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'])
+            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 expression.startswith("url://"):
+        #     url = expression.replace("url://", "")
+        #     url = url.format(
+        #         datasource_id=datasource.get("id"),
+        #         datasource_uid=datasource.get("uid"),
+        #         database_name=datasource.get("database"),
+        #     )
+        #     request_kwargs = {}
+        #     send_request = self.client.GET
+
+        # For all others, use the generic data source communication endpoint.
+        elif access_type in ["server", "proxy"]:
+            url = "/ds/query"
+            request_kwargs = {"json": request["data"]}
+
+        else:
+            raise NotImplementedError(f"Unable to submit query to data source with access type '{access_type}'")
+
+        # Submit query.
+        try:
+            r = send_request(url, **request_kwargs)
+            # logger.debug(f"Response from generic data source query: {r}")
+            return r
+        except (GrafanaClientError, GrafanaServerError) as ex:
+            logger.error(
+                f"Querying data source failed. id={datasource_id}, type={datasource_type}. "
+                f"Reason: {ex}. Response: {ex.response or '<empty>'}"
+            )
+            raise
 
     def health_check(self, datasource: Union[DatasourceIdentifier, Dict]) -> DatasourceHealthResponse:
         """