From 4ee3e45bf7a6808cc5b92d36aad1978c7f24955a Mon Sep 17 00:00:00 2001
From: Mar0xy <marie@kaifa.ch>
Date: Wed, 4 Oct 2023 02:24:20 +0200
Subject: [PATCH] upd: fix poll editing

---
 packages/backend/src/core/NoteEditService.ts  | 46 +++++++++++++++----
 .../frontend/src/scripts/use-note-capture.ts  |  9 +---
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/packages/backend/src/core/NoteEditService.ts b/packages/backend/src/core/NoteEditService.ts
index d38962fe3..8b26e6fb2 100644
--- a/packages/backend/src/core/NoteEditService.ts
+++ b/packages/backend/src/core/NoteEditService.ts
@@ -5,7 +5,7 @@
 
 import { setImmediate } from 'node:timers/promises';
 import * as mfm from 'mfm-js';
-import { In, LessThan } from 'typeorm';
+import { DataSource, In, LessThan } from 'typeorm';
 import * as Redis from 'ioredis';
 import { Inject, Injectable, OnApplicationShutdown } from '@nestjs/common';
 import RE2 from 're2';
@@ -20,7 +20,7 @@ import type { MiApp } from '@/models/App.js';
 import { concat } from '@/misc/prelude/array.js';
 import { IdService } from '@/core/IdService.js';
 import type { MiUser, MiLocalUser, MiRemoteUser } from '@/models/User.js';
-import type { IPoll } from '@/models/Poll.js';
+import { MiPoll, type IPoll } from '@/models/Poll.js';
 import { checkWordMute } from '@/misc/check-word-mute.js';
 import type { MiChannel } from '@/models/Channel.js';
 import { normalizeForSearch } from '@/misc/normalize-for-search.js';
@@ -154,6 +154,9 @@ export class NoteEditService implements OnApplicationShutdown {
 		@Inject(DI.redis)
 		private redisClient: Redis.Redis,
 
+		@Inject(DI.db)
+		private db: DataSource,
+
 		@Inject(DI.usersRepository)
 		private usersRepository: UsersRepository,
 
@@ -419,7 +422,27 @@ export class NoteEditService implements OnApplicationShutdown {
 			}));
 		}
 
-		await this.notesRepository.update(oldnote.id, note);
+		if (data.poll != null) {
+			// Start transaction
+			await this.db.transaction(async transactionalEntityManager => {
+				await transactionalEntityManager.update(MiNote, oldnote.id, note);
+
+				const poll = new MiPoll({
+					noteId: note.id,
+					choices: data.poll!.choices,
+					expiresAt: data.poll!.expiresAt,
+					multiple: data.poll!.multiple,
+					votes: new Array(data.poll!.choices.length).fill(0),
+					noteVisibility: note.visibility,
+					userId: user.id,
+					userHost: user.host,
+				});
+
+				await transactionalEntityManager.update(MiPoll, oldnote.id, poll);
+			});
+		} else {
+			await this.notesRepository.update(oldnote.id, note);
+		};
 
 		if (data.channel) {
 			this.redisClient.xadd(
@@ -483,6 +506,7 @@ export class NoteEditService implements OnApplicationShutdown {
 
 		if (data.poll && data.poll.expiresAt) {
 			const delay = data.poll.expiresAt.getTime() - Date.now();
+			this.queueService.endedPollNotificationQueue.remove(note.id);
 			this.queueService.endedPollNotificationQueue.add(note.id, {
 				noteId: note.id,
 			}, {
@@ -521,11 +545,17 @@ export class NoteEditService implements OnApplicationShutdown {
 
 			// Pack the note
 			const noteObj = await this.noteEntityService.pack(note);
-
-			this.globalEventService.publishNoteStream(note.id, 'updated', {
-				cw: note.cw,
-				text: note.text!,
-			});
+			if (data.poll != null) {
+				this.globalEventService.publishNoteStream(note.id, 'updated', {
+					cw: note.cw,
+					text: note.text!,
+				});
+			} else {
+				this.globalEventService.publishNoteStream(note.id, 'updated', {
+					cw: note.cw,
+					text: note.text!
+				});
+			}
 
 			this.roleService.addNoteToRoleTimeline(noteObj);
 
diff --git a/packages/frontend/src/scripts/use-note-capture.ts b/packages/frontend/src/scripts/use-note-capture.ts
index 42c1ebcca..c8bcb6d09 100644
--- a/packages/frontend/src/scripts/use-note-capture.ts
+++ b/packages/frontend/src/scripts/use-note-capture.ts
@@ -72,19 +72,12 @@ export function useNoteCapture(props: {
 				break;
 			}
 
-			case 'updated': {
-				note.value.updatedAt = new Date().toISOString();
-				note.value.cw = body.cw;
-				note.value.text = body.text;
-				break;
-			}
-
 			case 'deleted': {
 				props.isDeletedRef.value = true;
 				break;
 			}
 
-			case "updated": {
+			case 'updated': {
 				const editedNote = await os.api("notes/show", {
 					noteId: id,
 				});