From bbcce0d56f956d92a4487f7a9aa27516fe6b4c97 Mon Sep 17 00:00:00 2001
From: Karan Kajla <kkajla12@users.noreply.github.com>
Date: Sun, 31 Mar 2024 16:27:09 -0700
Subject: [PATCH] Add context support to query method

---
 src/modules/WarrantModule.ts | 22 +++++++++++++++++-----
 src/types/Query.ts           |  7 ++++++-
 src/types/index.ts           |  1 +
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/modules/WarrantModule.ts b/src/modules/WarrantModule.ts
index 6cba346..68c78b0 100644
--- a/src/modules/WarrantModule.ts
+++ b/src/modules/WarrantModule.ts
@@ -3,6 +3,7 @@ import { API_VERSION } from "../constants";
 import {
     ListResponse,
     ListWarrantParams,
+    QueryParams,
     QueryListParams,
     QueryResult,
     Warrant,
@@ -49,13 +50,24 @@ export default class WarrantModule {
         });
     }
 
-    public static async query(query: string, listParams: QueryListParams = {}, options: WarrantRequestOptions = {}): Promise<ListResponse<QueryResult>> {
+    public static async query(query: string | QueryParams, listParams: QueryListParams = {}, options: WarrantRequestOptions = {}): Promise<ListResponse<QueryResult>> {
+        const params: any = {
+            ...listParams,
+        };
+
+        // preserve backwards compatibility
+        if (typeof query === "string") {
+            params.q = query;
+        } else {
+            params.q = query.query;
+            if (query.context) {
+                params.context = query.context;
+            }
+        }
+
         return await WarrantClient.httpClient.get({
             url: `/${API_VERSION}/query`,
-            params: {
-                q: query,
-                ...listParams,
-            },
+            params: params,
             options,
         });
     }
diff --git a/src/types/Query.ts b/src/types/Query.ts
index 6090192..678823e 100644
--- a/src/types/Query.ts
+++ b/src/types/Query.ts
@@ -1,6 +1,11 @@
-import { Warrant } from "./Warrant";
+import { Warrant, PolicyContext } from "./Warrant";
 import { ListParams } from "./List";
 
+export interface QueryParams {
+    query: string;
+    context?: PolicyContext;
+}
+
 export interface QueryResult {
     objectType: string;
     objectId: string;
diff --git a/src/types/index.ts b/src/types/index.ts
index b5d0d71..64190b4 100644
--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -58,6 +58,7 @@ export {
     Subject,
 } from "./Warrant";
 export {
+    QueryParams,
     QueryResult,
     QueryListParams,
 } from "./Query";