From a1ab6b2d0c29cc58b1cee7d37108042f29335c09 Mon Sep 17 00:00:00 2001
From: Karan Kajla <kkajla12@users.noreply.github.com>
Date: Mon, 1 Apr 2024 13:14:32 -0700
Subject: [PATCH] Add context support to query method (#56)

---
 src/modules/WarrantModule.ts | 22 +++++++++++++++++-----
 src/types/Query.ts           |  7 ++++++-
 src/types/index.ts           |  1 +
 test/LiveTest.spec.ts        | 22 +++++++++++++++++++++-
 4 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/src/modules/WarrantModule.ts b/src/modules/WarrantModule.ts
index 6cba346..8ce1c53 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 = JSON.stringify(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";
diff --git a/test/LiveTest.spec.ts b/test/LiveTest.spec.ts
index 34f493d..07d83dd 100644
--- a/test/LiveTest.spec.ts
+++ b/test/LiveTest.spec.ts
@@ -933,6 +933,12 @@ describe.skip('Live Test', function () {
             relation: "member",
             subject: role1,
         });
+        await this.warrant.Warrant.create({
+            object: permission1,
+            relation: "member",
+            subject: role2,
+            policy: "tenantId == 123",
+        });
         await this.warrant.Role.assignRoleToUser(userA.userId, role1.roleId);
         await this.warrant.Role.assignRoleToUser(userB.userId, role2.roleId);
 
@@ -958,6 +964,20 @@ describe.skip('Live Test', function () {
         assert.strictEqual("role", resultSet.results[0].warrant.subject.objectType);
         assert.strictEqual("role1", resultSet.results[0].warrant.subject.objectId);
 
+        resultSet = await this.warrant.Warrant.query({
+            query: "select permission where user:userB is member",
+            context: {
+                tenantId: 123,
+            },
+        });
+        assert.strictEqual(3, resultSet.results.length);
+        assert.strictEqual("permission", resultSet.results[0].objectType);
+        assert.strictEqual("perm1", resultSet.results[0].objectId);
+        assert.strictEqual("permission", resultSet.results[1].objectType);
+        assert.strictEqual("perm2", resultSet.results[1].objectId);
+        assert.strictEqual("permission", resultSet.results[2].objectType);
+        assert.strictEqual("perm3", resultSet.results[2].objectId);
+
         let warrantToken = await this.warrant.Role.delete(role1.roleId);
         assert(warrantToken);
         warrantToken = await this.warrant.Role.delete(role2.roleId);
@@ -972,5 +992,5 @@ describe.skip('Live Test', function () {
         assert(warrantToken);
         warrantToken = await this.warrant.User.delete(userB.userId);
         assert(warrantToken);
-    })
+    });
 });