Use PartitonKey In update/Delete queries (#60)

* updating to use requestOptions partition key instead of id as partition key
* added unit testing for partitionKey
* upgraded peerDependencies: @azure/cosmos

Co-authored-by: Andrej Kyselica <[email protected]>

Author: aspladypfl

package-lock.json

@@ -41,7 +41,7 @@
     "ync": "0.0.9"
   },
   "peerDependencies": {
-    "@azure/cosmos": "3.10.3"
+    "@azure/cosmos": "^3.10.3"
   },
   "dependencies": {
     "apollo-datasource": "0.8.0",

package.json

@@ -76,30 +76,30 @@ export class CosmosDataSource<TData extends { id: string }, TContext>
     return response;
   }
 
-  async deleteOne(id: string) {
+  async deleteOne(id: string, partitionKey?: string) {
     this.options?.logger?.info(
       `CosmosDataSource/deleteOne: deleting id: '${id}'`
     );
-    const response = await this.container.item(id, id).delete<TData>();
+    const response = await this.container.item(id, partitionKey).delete<TData>();
     await this.deleteFromCacheById(id);
     return response;
   }
 
-  async updateOne(updDoc: TData) {
+  async updateOne(updDoc: TData, partitionKey?: string) {
     const response = await this.container
-      .item(updDoc.id, updDoc.id)
+      .item(updDoc.id, partitionKey)
       .replace(updDoc);
     if (response.resource) {
       this.primeLoader(response.resource);
     }
     return response;
   }
 
-  async updateOnePartial(id: string, contents: Partial<TData>) {
+  async updateOnePartial(id: string, contents: Partial<TData>, partitionKey?: string) {
     this.options?.logger?.debug(
       `Updating doc id ${id} contents: ${JSON.stringify(contents, null, "")}`
     );
-    const item = this.container.item(id, id);
+    const item = this.container.item(id, partitionKey);
     const docItem = await item.read<TData>();
     const { resource } = docItem;
     const newResource = { ...resource, ...contents, id } as TData; // don't change the ID ever

src/datasource.ts

@@ -1,4 +1,4 @@
-import withCosmosDb from "./withCosmosDb";
+import withCosmosDb, {CosmosDB} from "./withCosmosDb";
 import { expect } from "chai";
 
 import { CosmosDataSource } from "../src/datasource";
@@ -8,6 +8,7 @@ process.env.NODE_TLS_REJECT_UNAUTHORIZED = "0";
 interface UserDoc {
   id: string;
   email: string;
+  partitionKey?: string;
 }
 
 interface Context {
@@ -67,3 +68,191 @@ describe("basic crud", () => {
     })
   );
 });
+
+describe("partitionKey crud tests",function() {
+    let cosmosDB: CosmosDB
+
+    before(async function () {
+        cosmosDB = new CosmosDB();
+        await cosmosDB.createDatabase();
+    })
+
+    after(function() {
+      cosmosDB.close();
+    })
+    describe("with a partitionKey", function() {
+      let userDataSource: UserDataSource
+      
+      beforeEach(async function () {     
+        const { container } = await cosmosDB.createContainer({
+          id: "test-collection",
+          partitionKey: "/partitionKey",
+        });
+
+        userDataSource = new UserDataSource(container);
+        userDataSource.initialize({});
+      })
+
+      afterEach(async function () {    
+        userDataSource.container.delete();
+      })
+      
+      it("should create and find a user", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+          partitionKey: "examplePartitionKey",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+
+        // read the user by ID
+        const user1Find1Resp = await userDataSource.findOneById(user1.id);
+        expect(user1Find1Resp?.email).to.equal(user1.email);
+        expect(user1Find1Resp?.id).to.equal(user1.id);
+      })
+
+      it("should create and update a user and find with query", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+          partitionKey: "examplePartitionKey",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+
+        // update the user
+        const newEmail = "[email protected]";
+        await userDataSource.updateOne({
+          ...user1,
+          email: newEmail,
+        });
+
+        // read the user with a query
+        const user1FindQueryResp = await userDataSource.findManyByQuery(
+          `select * from c where c.id = '${user1.id}'`
+        );
+        expect(user1FindQueryResp.resources.length).to.equal(1);
+        expect(user1FindQueryResp.resources[0].email).to.equal(newEmail);
+      })
+      
+      it("should create and delete user", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+          partitionKey: "examplePartitionKey",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+
+        // delete the user
+        await userDataSource.deleteOne(user1.id, user1.partitionKey);
+
+        // try to read the user back out (shouldn't exist)
+        const user1Find1Resp2 = await userDataSource.findOneById(user1.id);
+        expect(user1Find1Resp2).to.equal(undefined);
+      })
+      
+      it("should create and not delete user because partitionKey isn't provided", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+          partitionKey: "examplePartitionKey",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+        
+        // delete the user will fail because partitionKey isn't provided
+        userDataSource.deleteOne(user1.id);
+
+        const user1Find1Resp2 = await userDataSource.findOneById(user1.id);
+
+        expect(user1Find1Resp2.id).to.equal(user1.id);
+        expect(user1Find1Resp2.email).to.equal(user1.email);
+        expect(user1Find1Resp2.partitionKey).to.equal(user1.partitionKey);
+      })
+    })
+    describe("without a partitionKey", function() {
+      let userDataSource: UserDataSource
+      
+      beforeEach(async function () {     
+        const { container } = await cosmosDB.createContainer({
+          id: "test-collection",
+        });
+
+        userDataSource = new UserDataSource(container);
+        userDataSource.initialize({});
+      })
+
+      afterEach(async function () {    
+        userDataSource.container.delete();
+      })
+      
+      it("should create and find a user", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+
+        // read the user by ID
+        const user1Find1Resp = await userDataSource.findOneById(user1.id);
+        expect(user1Find1Resp?.email).to.equal(user1.email);
+        expect(user1Find1Resp?.id).to.equal(user1.id);
+      })
+
+      it("should create and update a user and find with query", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+
+        // update the user
+        const newEmail = "[email protected]";
+        await userDataSource.updateOne({
+          ...user1,
+          email: newEmail,
+        });
+
+        // read the user with a query
+        const user1FindQueryResp = await userDataSource.findManyByQuery(
+          `select * from c where c.id = '${user1.id}'`
+        );
+        expect(user1FindQueryResp.resources.length).to.equal(1);
+        expect(user1FindQueryResp.resources[0].email).to.equal(newEmail);
+      })
+      
+      it("should create and delete user", async function() { 
+        const user1 = {
+          id: "us_one",
+          email: "[email protected]",
+        };
+
+        // create the user
+        const user1Resp = await userDataSource.createOne(user1);
+        expect(user1Resp.resource?.id).to.equal("us_one");
+
+        // delete the user
+        await userDataSource.deleteOne(user1.id);
+
+        // try to read the user back out (shouldn't exist)
+        const user1Find1Resp2 = await userDataSource.findOneById(user1.id);
+        expect(user1Find1Resp2).to.equal(undefined);
+      })
+    })
+})

tests/integration.test.ts

@@ -1,4 +1,4 @@
-import { CosmosClient } from "@azure/cosmos";
+import { CosmosClient, Database, ContainerRequest  } from "@azure/cosmos";
 import * as net from "net";
 import cosmosDBServer from "@zeit/cosmosdb-server";
 
@@ -27,3 +27,32 @@ export default function withCosmosDBServer<R, T extends []>(
     }
   };
 }
+
+export class CosmosDB {  
+  #server: net.Server
+  #database: Database
+  #client: CosmosClient
+
+  constructor() {
+    this.#server = cosmosDBServer();
+    this.#server.listen(0)
+    const { port } = this.#server.address() as net.AddressInfo;
+    this.#client = new CosmosClient({
+      endpoint: `https://localhost:${port}`,
+      key: "test-master-key",
+    });
+  }
+  async createDatabase(dbId = "test-database") {
+    this.#database = (await this.#client.databases.create({
+      id: dbId,
+    })).database;
+  }
+  async createContainer(containerRequest: ContainerRequest) {
+    return await this.#database.containers.create(
+      containerRequest
+    );
+  }
+  close() {
+    this.#server.close();
+  }
+}

tests/withCosmosDb.ts

@@ -3957,4 +3957,4 @@ [email protected]:
 yocto-queue@^0.1.0:
   version "0.1.0"
   resolved "https://registry.npmjs.org/yocto-queue/-/yocto-queue-0.1.0.tgz"
-  integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==
+  integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==