Add timeout parameter to query (#32)

Author: haohanyang

pkg/plugin/datasource.go

@@ -105,7 +105,13 @@ func (d *Datasource) query(ctx context.Context, _ backend.PluginContext, query b
 		return backend.ErrDataResponse(backend.StatusBadRequest, fmt.Sprintf("Failed to unmarshal JsonExt: %v", err.Error()))
 	}
 
-	cursor, err := db.Collection(qm.Collection).Aggregate(ctx, pipeline)
+	aggregateOpts := options.AggregateOptions{}
+	if qm.Timeout > 0 {
+		aggregateOpts.MaxTime = pointer(time.Hour * time.Duration(qm.Timeout))
+		backend.Logger.Debug("Aggregate timeout was set", "timeout", qm.Timeout)
+	}
+
+	cursor, err := db.Collection(qm.Collection).Aggregate(ctx, pipeline, &aggregateOpts)
 	if err != nil {
 		return backend.ErrDataResponse(backend.StatusBadRequest, fmt.Sprintf("Failed to query: %v", err.Error()))
 

pkg/plugin/types.go

@@ -18,6 +18,7 @@ type queryModel struct {
 	Collection    string `json:"collection"`
 	QueryType     string `json:"queryType"`
 	QueryLanguage string `json:"queryLanguage"`
+	Timeout       int    `json:"timeout"`
 }
 
 type TimeSeriesRow[T any] struct {

src/components/QueryEditor.tsx

@@ -14,7 +14,7 @@ import {
 import { QueryEditorProps, SelectableValue } from "@grafana/data";
 import { DataSource } from "../datasource";
 import { MongoDataSourceOptions, MongoQuery, QueryLanguage, QueryType, DEFAULT_QUERY } from "../types";
-import { parseJsQuery, parseJsQueryLegacy, validateJsonQueryText } from "../utils";
+import { parseJsQuery, parseJsQueryLegacy, validateJsonQueryText, validateTimeout } from "../utils";
 import * as monacoType from "monaco-editor/esm/vs/editor/editor.api";
 
 type Props = QueryEditorProps<DataSource, MongoQuery, MongoDataSourceOptions>;
@@ -35,6 +35,7 @@ export function QueryEditor({ query, onChange, onRunQuery }: Props) {
 
   const codeEditorRef = useRef<monacoType.editor.IStandaloneCodeEditor | null>(null);
   const [queryTextError, setQueryTextError] = useState<string | null>(null);
+  const [timeoutText, setTimeoutText] = useState<string>(query.timeout ? query.timeout.toString() : "");
 
   const optionsLanguage = [
     { label: "JSON", value: QueryLanguage.JSON },
@@ -74,6 +75,16 @@ export function QueryEditor({ query, onChange, onRunQuery }: Props) {
     onChange({ ...query, collection: event.target.value });
   };
 
+  const onTimeoutChange = (event: ChangeEvent<HTMLInputElement>) => {
+    setTimeoutText(event.target.value);
+    console.log(event.target.value);
+    if (!event.target.value) {
+      onChange({ ...query, timeout: undefined });
+    } else if (validateTimeout(event.target.value)) {
+      onChange({ ...query, timeout: parseInt(event.target.value, 10) });
+    }
+  };
+
   const onCodeEditorDidMount = (e: monacoType.editor.IStandaloneCodeEditor) => {
     codeEditorRef.current = e;
   };
@@ -98,6 +109,10 @@ export function QueryEditor({ query, onChange, onRunQuery }: Props) {
           error="Please enter the collection" invalid={!query.collection}>
           <Input id="query-editor-collection" onChange={onCollectionChange} value={query.collection} required />
         </InlineField>
+        <InlineField label="Timeout" tooltip="(Optional) The maximum amount of time (in milisecond) that the query can run on the server."
+          error="Invalid timeout" invalid={timeoutText !== "" && !validateTimeout(timeoutText)}>
+          <Input id="query-editor-timeout" onChange={onTimeoutChange} value={timeoutText} />
+        </InlineField>
       </InlineFieldRow>
       <Divider />
       <InlineField label="Query language">

src/types.ts

@@ -6,6 +6,7 @@ export interface MongoQuery extends DataQuery {
   collection?: string;
   queryType?: string;
   queryLanguage?: string;
+  timeout?: number;
 }
 
 export const QueryType = {

src/utils.ts

@@ -121,3 +121,13 @@ export function getMetricValues(response: DataQueryResponse): MetricFindValue[]
         };
     });
 }
+
+export function validateTimeout(timeout: string) {
+
+    if (!/^\d+$/.test(timeout.trim())) {
+        return false;
+    }
+
+    const parsed = parseInt(timeout, 10);
+    return parsed > 0;
+}