From 509919da84a8acc07c4c47fff83ad1a720273317 Mon Sep 17 00:00:00 2001
From: Jay Lee <jay0lee@gmail.com>
Date: Wed, 14 Sep 2022 18:55:17 +0000
Subject: [PATCH] Reduce memory with shelve. Fixes #1560

---
 src/gam/__init__.py      |  1 +
 src/gam/gapi/__init__.py | 29 +++++++++++++++++++++++++----
 src/gam/var.py           |  7 +++++++
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/src/gam/__init__.py b/src/gam/__init__.py
index da8014cd..d8df7e38 100755
--- a/src/gam/__init__.py
+++ b/src/gam/__init__.py
@@ -553,6 +553,7 @@ def SetGlobalVariables():
                       'debug.gam',
                       filePresentValue=4,
                       fileAbsentValue=0)
+    _getOldSignalFile(GC_LOW_MEMORY, 'lowmemory.txt')
     _getOldSignalFile(GC_NO_BROWSER, 'nobrowser.txt')
     _getOldSignalFile(GC_NO_TDEMAIL, 'notdemail.txt')
     _getOldSignalFile(GC_OAUTH_BROWSER, 'oauthbrowser.txt')
diff --git a/src/gam/gapi/__init__.py b/src/gam/gapi/__init__.py
index 2bd6bdde..160d5a5d 100644
--- a/src/gam/gapi/__init__.py
+++ b/src/gam/gapi/__init__.py
@@ -1,6 +1,9 @@
 """Methods related to execution of GAPI requests."""
 
+import os.path
+import shelve
 import sys
+from tempfile import TemporaryDirectory
 
 import googleapiclient.errors
 import google.auth.exceptions
@@ -10,7 +13,8 @@ from gam import controlflow
 from gam import display
 from gam.gapi import errors
 from gam import transport
-from gam.var import (GM_Globals, GM_CURRENT_API_SCOPES, GM_CURRENT_API_USER,
+from gam.var import (GC_Values, GC_LOW_MEMORY, GM_Globals,
+                     GM_CURRENT_API_SCOPES, GM_CURRENT_API_USER,
                      GM_EXTRA_ARGS_DICT, GM_OAUTH2SERVICE_ACCOUNT_CLIENT_ID,
                      MAX_RESULTS_API_EXCEPTIONS, MESSAGE_API_ACCESS_CONFIG,
                      MESSAGE_API_ACCESS_DENIED, MESSAGE_SERVICE_NOT_APPLICABLE)
@@ -238,8 +242,13 @@ def process_page(page, items, all_items, total_items, page_message, message_attr
         page_items = page.get(items, [])
         num_page_items = len(page_items)
         total_items += num_page_items
-        if all_items is not None:
+        if type(all_items) is list:
             all_items.extend(page_items)
+        elif all_items is not None:
+            i = len(all_items)
+            for item in page_items:
+                all_items[str(i)] = item
+                i += 1
     else:
         page_token = None
         num_page_items = 0
@@ -273,6 +282,7 @@ def finalize_page_message(page_message):
         sys.stderr.write('\r\n')
         sys.stderr.flush()
 
+
 def get_all_pages(service,
                   function,
                   items='items',
@@ -328,7 +338,17 @@ def get_all_pages(service,
                 kwargs['body'].update(page_key)
             else:
                 kwargs.update(page_key)
-    all_items = []
+    if GC_Values[GC_LOW_MEMORY]:
+        tempdir = TemporaryDirectory(
+                prefix='GAM-',
+                ignore_cleanup_errors=True)
+        tempfile = os.path.join(tempdir.name, 'gapi_pages')
+        all_items = shelve.open(tempfile)
+        # attach tempdir to all_items so we
+        # don't cleanup tempdir early
+        all_items._tempdir = tempdir
+    else:
+        all_items = []
     page_token = None
     total_items = 0
     while True:
@@ -341,13 +361,14 @@ def get_all_pages(service,
         page_token, total_items = process_page(page, items, all_items, total_items, page_message, message_attribute)
         if not page_token:
             finalize_page_message(page_message)
+            if type(all_items) is not list:
+                all_items = all_items.values()
             return all_items
         if page_args_in_body:
             kwargs['body']['pageToken'] = page_token
         else:
             kwargs['pageToken'] = page_token
 
-
 # TODO: Make this private once all execution related items that use this method
 # have been brought into this file
 def handle_oauth_token_error(e, soft_errors):
diff --git a/src/gam/var.py b/src/gam/var.py
index c481f358..b9f327bf 100644
--- a/src/gam/var.py
+++ b/src/gam/var.py
@@ -1270,6 +1270,9 @@ GC_ENABLE_DASA = 'enabledasa'
 # and doRequestOAuth prints a link and waits for the verification code when
 # oauth2.txt is being created
 GC_NO_BROWSER = 'no_browser'
+# If low memory is True, GAM tries to save RAM by writing pages to disk
+# temporarily
+GC_LOW_MEMORY = 'low_memory'
 # If no_tdemail is True, writeCSVfile won't send an email
 GC_NO_TDEMAIL = 'no_tdemail'
 # oauth_browser forces usage of web server OAuth flow that proved problematic.
@@ -1325,6 +1328,7 @@ GC_Defaults = {
     GC_DOMAIN: '',
     GC_DRIVE_DIR: '',
     GC_ENABLE_DASA: False,
+    GC_LOW_MEMORY: False,
     GC_NO_BROWSER: False,
     GC_NO_TDEMAIL: False,
     GC_NO_CACHE: False,
@@ -1409,6 +1413,9 @@ GC_VAR_INFO = {
     GC_ENABLE_DASA: {
         GC_VAR_TYPE: GC_TYPE_BOOLEAN
     },
+    GC_LOW_MEMORY: {
+        GC_VAR_TYPE: GC_TYPE_BOOLEAN
+    },
     GC_NO_BROWSER: {
         GC_VAR_TYPE: GC_TYPE_BOOLEAN
     },