Improve upload account balance history

README.md

@@ -134,7 +134,7 @@ As of writing this README, the following methods are supported:
 - `set_budget_amount` - sets a budget's value to the given amount (date allowed, will only apply to month specified by default). A zero amount value will "unset" or "clear" the budget for the given category.
 - `create_manual_account` - creates a new manual account
 - `delete_account` - deletes an account by the provided account id
-- `upload_account_balance_history` - uploads account history csv file for a given account
+- `upload_account_balance_history` - uploads and parses account history csv file for a given account
 
 # Contributing
 

monarchmoney/monarchmoney.py

@@ -18,9 +18,11 @@
 AUTH_HEADER_KEY = "authorization"
 CSRF_KEY = "csrftoken"
 DEFAULT_RECORD_LIMIT = 100
+DELAY = 10
 ERRORS_KEY = "error_code"
 SESSION_DIR = ".mm"
 SESSION_FILE = f"{SESSION_DIR}/mm_session.pickle"
+TIMEOUT = 300
 
 
 class MonarchMoneyEndpoints(object):
@@ -127,6 +129,17 @@ async def multi_factor_authenticate(
         """Performs multi-factor authentication to access a Monarch Money account."""
         await self._multi_factor_authenticate(email, password, code)
 
+    async def _upload_form_data(self, url: str, data: FormData) -> dict:
+        """
+        Retrieves the response from the server for a given URL and form data.
+        """
+        async with ClientSession(headers=self._headers) as session:
+            resp = await session.post(url, data=data)
+            if resp.status != 200:
+                raise RequestFailedException(f"HTTP Code {resp.status}: {resp.reason}")
+
+            return await resp.json()
+
     async def get_accounts(self) -> Dict[str, Any]:
         """
         Gets the list of accounts configured in the Monarch Money account.
@@ -434,8 +447,8 @@ async def is_accounts_refresh_complete(self) -> bool:
     async def request_accounts_refresh_and_wait(
         self,
         account_ids: Optional[List[str]] = None,
-        timeout: int = 300,
-        delay: int = 10,
+        timeout: int = TIMEOUT,
+        delay: int = DELAY,
     ) -> bool:
         """
         Convenience method for forcing an accounts refresh on Monarch, as well
@@ -2394,13 +2407,20 @@ async def set_budget_amount(
         )
 
     async def upload_account_balance_history(
-        self, account_id: str, csv_content: str
-    ) -> None:
+        self,
+        account_id: str,
+        csv_content: str,
+        timeout: int = TIMEOUT,
+        delay: int = DELAY,
+    ) -> str:
         """
-        Uploads the account balance history csv for a given account.
+        Uploads the account balance history CSV for a specified account.
 
         :param account_id: The account ID to apply the history to.
         :param csv_content: CSV representation of the balance history.
+                            Headers: Date, Amount, and Account Name.
+        :param timeout: The number of seconds to wait before timing out
+        :param delay: The number of seconds to wait for each check on whether parsing is completed
         """
         if not account_id or not csv_content:
             raise RequestFailedException("account_id and csv_content cannot be empty")
@@ -2410,13 +2430,94 @@ async def upload_account_balance_history(
         form.add_field("files", csv_content, filename=filename, content_type="text/csv")
         form.add_field("account_files_mapping", json.dumps({filename: account_id}))
 
-        async with ClientSession(headers=self._headers) as session:
-            resp = await session.post(
-                MonarchMoneyEndpoints.getAccountBalanceHistoryUploadEndpoint(),
-                data=form,
-            )
-            if resp.status != 200:
-                raise RequestFailedException(f"HTTP Code {resp.status}: {resp.reason}")
+        upload_response = await self._upload_form_data(
+            url=MonarchMoneyEndpoints.getAccountBalanceHistoryUploadEndpoint(),
+            data=form,
+        )
+
+        session_key = upload_response["session_key"]
+
+        parse_response = await self._initiate_upload_balance_history_session(
+            session_key=session_key
+        )
+
+        is_completed = (
+            parse_response["parseBalanceHistory"]["uploadBalanceHistorySession"][
+                "status"
+            ]
+            == "completed"
+        )
+
+        start = time.time()
+        while not is_completed and (time.time() <= (start + timeout)):
+            await asyncio.sleep(delay)
+
+            is_completed = (
+                await self._is_upload_balance_history_complete(session_key)
+            )["uploadBalanceHistorySession"]["status"] == "completed"
+
+        return is_completed
+
+    async def _initiate_upload_balance_history_session(self, session_key: str) -> dict:
+        """
+        Triggers parsing of the uploaded balance history CSV file.
+
+        :param session_key: The session key for the uploaded file.
+        """
+
+        query = gql(
+            """
+            mutation Web_ParseUploadBalanceHistorySession($input: ParseBalanceHistoryInput!) {
+                parseBalanceHistory(input: $input) {
+                    uploadBalanceHistorySession {
+                        ...UploadBalanceHistorySessionFields
+                        __typename
+                    }
+                    __typename
+                }
+            }
+            fragment UploadBalanceHistorySessionFields on UploadBalanceHistorySession {
+                sessionKey
+                status
+                __typename
+            }
+            """
+        )
+
+        variables = {"input": {"sessionKey": session_key}}
+
+        return await self.gql_call(
+            "Web_ParseUploadBalanceHistorySession", query, variables
+        )
+
+    async def _is_upload_balance_history_complete(self, session_key: str):
+        """
+        Retrieves the status of the upload balance history session.
+
+        :param session_key: The session key for the uploaded file.
+        """
+
+        query = gql(
+            """
+            query Web_GetUploadBalanceHistorySession($sessionKey: String!) {
+                uploadBalanceHistorySession(sessionKey: $sessionKey) {
+                    ...UploadBalanceHistorySessionFields
+                    __typename
+                }
+            }
+            fragment UploadBalanceHistorySessionFields on UploadBalanceHistorySession {
+                sessionKey
+                status
+                __typename
+            }
+            """
+        )
+
+        variables = {"sessionKey": session_key}
+
+        return await self.gql_call(
+            "Web_GetUploadBalanceHistorySession", query, variables
+        )
 
     async def get_recurring_transactions(
         self,