From e5a8247aa344f6efccba83938effa1930ac9d5e9 Mon Sep 17 00:00:00 2001
From: DJ Mountney <dmountney@pop-os.localdomain>
Date: Fri, 10 Nov 2023 15:38:09 -0800
Subject: [PATCH] Learn Categories added to ynab5 import

---
 packages/api/methods.js                       | 13 ++++++++--
 .../loot-core/src/server/accounts/sync.ts     | 10 +++++---
 .../src/server/accounts/transactions.ts       | 25 +++++++++++--------
 packages/loot-core/src/server/api.ts          |  7 +++++-
 .../loot-core/src/server/importers/ynab5.ts   |  4 ++-
 5 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/packages/api/methods.js b/packages/api/methods.js
index dad30f1ef54..2b4f497380b 100644
--- a/packages/api/methods.js
+++ b/packages/api/methods.js
@@ -58,8 +58,17 @@ export function setBudgetCarryover(month, categoryId, flag) {
   return send('api/budget-set-carryover', { month, categoryId, flag });
 }
 
-export function addTransactions(accountId, transactions) {
-  return send('api/transactions-add', { accountId, transactions });
+export function addTransactions(
+  accountId,
+  transactions,
+  { learnCategories = false, runTransfers = false } = {},
+) {
+  return send('api/transactions-add', {
+    accountId,
+    transactions,
+    learnCategories: learnCategories,
+    runTransfers: runTransfers,
+  });
 }
 
 export function importTransactions(accountId, transactions) {
diff --git a/packages/loot-core/src/server/accounts/sync.ts b/packages/loot-core/src/server/accounts/sync.ts
index ae80d775162..7dba12fc328 100644
--- a/packages/loot-core/src/server/accounts/sync.ts
+++ b/packages/loot-core/src/server/accounts/sync.ts
@@ -705,7 +705,7 @@ export async function reconcileTransactions(acctId, transactions) {
 export async function addTransactions(
   acctId,
   transactions,
-  { runTransfers = true } = {},
+  { runTransfers = true, learnCategories = false } = {},
 ) {
   const added = [];
 
@@ -737,8 +737,12 @@ export async function addTransactions(
   await createNewPayees(payeesToCreate, added);
 
   let newTransactions;
-  if (runTransfers) {
-    let res = await batchUpdateTransactions({ added });
+  if (runTransfers || learnCategories) {
+    let res = await batchUpdateTransactions({
+      added,
+      learnCategories: learnCategories,
+      runTransfers: runTransfers,
+    });
     newTransactions = res.added.map(t => t.id);
   } else {
     await batchMessages(async () => {
diff --git a/packages/loot-core/src/server/accounts/transactions.ts b/packages/loot-core/src/server/accounts/transactions.ts
index ed1591290e9..30f4ea94ae9 100644
--- a/packages/loot-core/src/server/accounts/transactions.ts
+++ b/packages/loot-core/src/server/accounts/transactions.ts
@@ -40,6 +40,7 @@ export async function batchUpdateTransactions({
   updated,
   learnCategories = false,
   detectOrphanPayees = true,
+  runTransfers = true,
 }: {
   added?: Array<{ id: string; payee: unknown; category: unknown }>;
   deleted?: Array<{ id: string; payee: unknown }>;
@@ -51,6 +52,7 @@ export async function batchUpdateTransactions({
   }>;
   learnCategories?: boolean;
   detectOrphanPayees?: boolean;
+  runTransfers?: boolean;
 }) {
   // Track the ids of each type of transaction change (see below for why)
   let addedIds = [];
@@ -126,18 +128,21 @@ export async function batchUpdateTransactions({
   // to the client so that can apply them. Note that added
   // transactions just return the full transaction.
   let resultAdded = allAdded;
-  let resultUpdated: Awaited<ReturnType<typeof transfer.onUpdate>>[];
+  let resultUpdated = allUpdated;
+  let transfersUpdated: Awaited<ReturnType<typeof transfer.onUpdate>>[];
 
-  await batchMessages(async () => {
-    await Promise.all(allAdded.map(t => transfer.onInsert(t)));
+  if (runTransfers) {
+    await batchMessages(async () => {
+      await Promise.all(allAdded.map(t => transfer.onInsert(t)));
 
-    // Return any updates from here
-    resultUpdated = (
-      await Promise.all(allUpdated.map(t => transfer.onUpdate(t)))
-    ).filter(Boolean);
+      // Return any updates from here
+      transfersUpdated = (
+        await Promise.all(allUpdated.map(t => transfer.onUpdate(t)))
+      ).filter(Boolean);
 
-    await Promise.all(allDeleted.map(t => transfer.onDelete(t)));
-  });
+      await Promise.all(allDeleted.map(t => transfer.onDelete(t)));
+    });
+  }
 
   if (learnCategories) {
     // Analyze any updated categories and update rules to learn from
@@ -176,6 +181,6 @@ export async function batchUpdateTransactions({
 
   return {
     added: resultAdded,
-    updated: resultUpdated,
+    updated: runTransfers ? transfersUpdated : resultUpdated,
   };
 }
diff --git a/packages/loot-core/src/server/api.ts b/packages/loot-core/src/server/api.ts
index 01190365df1..50d4cb97a75 100644
--- a/packages/loot-core/src/server/api.ts
+++ b/packages/loot-core/src/server/api.ts
@@ -393,9 +393,14 @@ handlers['api/transactions-import'] = withMutation(async function ({
 handlers['api/transactions-add'] = withMutation(async function ({
   accountId,
   transactions,
+  runTransfers = false,
+  learnCategories = false,
 }) {
   checkFileOpen();
-  await addTransactions(accountId, transactions, { runTransfers: false });
+  await addTransactions(accountId, transactions, {
+    runTransfers: runTransfers,
+    learnCategories: learnCategories,
+  });
   return 'ok';
 });
 
diff --git a/packages/loot-core/src/server/importers/ynab5.ts b/packages/loot-core/src/server/importers/ynab5.ts
index 8641d82cb22..6268ffba4aa 100644
--- a/packages/loot-core/src/server/importers/ynab5.ts
+++ b/packages/loot-core/src/server/importers/ynab5.ts
@@ -234,7 +234,9 @@ async function importTransactions(
         })
         .filter(x => x);
 
-      await actual.addTransactions(entityIdMap.get(accountId), toImport);
+      await actual.addTransactions(entityIdMap.get(accountId), toImport, {
+        learnCategories: true,
+      });
     }),
   );
 }