diff --git a/server/presenters/policy.ts b/server/presenters/policy.ts
index 2c8670223..472b7975c 100644
--- a/server/presenters/policy.ts
+++ b/server/presenters/policy.ts
@@ -1,3 +1,4 @@
+import { compact } from "lodash";
 import { traceFunction } from "@server/logging/tracing";
 import { User } from "@server/models";
 
@@ -6,11 +7,14 @@ type Policy = {
   abilities: Record<string, boolean>;
 };
 
-function presentPolicy(user: User, objects: Record<string, any>[]): Policy[] {
+function presentPolicy(
+  user: User,
+  objects: (Record<string, any> | null)[]
+): Policy[] {
   // eslint-disable-next-line @typescript-eslint/no-var-requires
   const { serialize } = require("../policies");
 
-  return objects.map((object) => ({
+  return compact(objects).map((object) => ({
     id: object.id,
     abilities: serialize(user, object),
   }));
diff --git a/server/routes/api/documents/documents.test.ts b/server/routes/api/documents/documents.test.ts
index 796ade083..3e580ca9f 100644
--- a/server/routes/api/documents/documents.test.ts
+++ b/server/routes/api/documents/documents.test.ts
@@ -2818,6 +2818,39 @@ describe("#documents.update", () => {
     expect(res.status).toBe(400);
     expect(body.message).toBe("id: Required");
   });
+
+  describe("apiVersion=2", () => {
+    it("should successfully publish a draft", async () => {
+      const { user, team, collection } = await seed();
+      const document = await buildDraftDocument({
+        title: "title",
+        text: "text",
+        teamId: team.id,
+      });
+
+      const res = await server.post("/api/documents.update", {
+        body: {
+          apiVersion: 2,
+          token: user.getJwtToken(),
+          id: document.id,
+          title: "Updated title",
+          text: "Updated text",
+          lastRevision: document.revisionCount,
+          collectionId: collection.id,
+          publish: true,
+        },
+      });
+      const body = await res.json();
+      expect(res.status).toEqual(200);
+      expect(body.data.document.collectionId).toBe(collection.id);
+      expect(body.data.document.title).toBe("Updated title");
+      expect(body.data.document.text).toBe("Updated text");
+      expect(body.data.collection.icon).toBe(collection.icon);
+      expect(body.data.collection.documents.length).toBe(
+        collection.documentStructure!.length + 1
+      );
+    });
+  });
 });
 
 describe("#documents.archive", () => {
diff --git a/server/routes/api/documents/documents.ts b/server/routes/api/documents/documents.ts
index 82d5d8cd3..91fad69a9 100644
--- a/server/routes/api/documents/documents.ts
+++ b/server/routes/api/documents/documents.ts
@@ -854,8 +854,8 @@ router.post(
       throw InvalidRequestError("Document has changed since last revision");
     }
 
-    await sequelize.transaction((transaction) => {
-      return documentUpdater({
+    collection = await sequelize.transaction(async (transaction) => {
+      await documentUpdater({
         document,
         user,
         title,
@@ -869,29 +869,26 @@ router.post(
         transaction,
         ip: ctx.request.ip,
       });
+
+      return await Collection.scope({
+        method: ["withMembership", user.id],
+      }).findByPk(document.collectionId, { transaction });
     });
 
     document.updatedBy = user;
-
-    // Original collection has membership data for calculating policies so we
-    // use that collection and just update the documentStructure for performance
-    if (collection) {
-      const { documentStructure } = collection;
-      document.collection = collection;
-      document.collection.documentStructure = documentStructure;
-    }
+    document.collection = collection;
 
     ctx.body = {
       data:
         apiVersion === 2
           ? {
               document: await presentDocument(document),
-              collection: document.collection
-                ? presentCollection(document.collection)
+              collection: collection
+                ? presentCollection(collection)
                 : undefined,
             }
           : await presentDocument(document),
-      policies: presentPolicies(user, [document]),
+      policies: presentPolicies(user, [document, collection]),
     };
   }
 );