Added pagination system to elude memory issues

pymongoexport_csv.py

@@ -7,28 +7,43 @@ import argparse
 from lib import utils
 
 
-# Command line parsing
-parser = argparse.ArgumentParser(
+def get_tweets_page(collection, page_size: int, num_page: int):
+    tweets_page = []
+    tweets = collection.find().skip(num_page * page_size).limit(page_size)
+    for tweet in tweets:
+        tweets_page.append(tweet)
+    return tweets_page
+
+
+if __name__ == "__main__":
+
+    # Command line parsing
+    parser = argparse.ArgumentParser(
         description="Dump the tweets of a database to a JSON file")
-parser.add_argument("-H", "--host", type=str, default="localhost")
-parser.add_argument("-p", "--port", type=int, default=27017)
-parser.add_argument("database", type=str)
-args = parser.parse_args()
+    parser.add_argument("-H", "--host", type=str, default="localhost")
+    parser.add_argument("-p", "--port", type=int, default=27017)
+    parser.add_argument("-s", "--pagesize", type=int, default=1000)
+    parser.add_argument("database", type=str)
+    args = parser.parse_args()
 
-# Dirs and files
-script_dir = os.path.dirname(__file__)
-output_dir = os.path.join(script_dir, "pymongodump", args.database)
-header_file = os.path.join(script_dir, "header.txt")
+    # Dirs and files
+    script_dir = os.path.dirname(__file__)
+    output_dir = os.path.join(script_dir, "pymongodump", args.database)
+    header_file = os.path.join(script_dir, "header.txt")
 
-# MongoDB connection
-client = pymongo.MongoClient(args.host, args.port)
-database_tweets = client[args.database]["tweets"]
+    # MongoDB connection
+    client = pymongo.MongoClient(args.host, args.port)
+    database_tweets = client[args.database]["tweets"]
 
-with open(header_file) as f:
+    with open(header_file) as f:
         header = f.readline()
+    buffer_tweets = {}
+    num_page = 0
 
-buffer_tweets = {}
-for tweet in database_tweets.find():
+    tweets_page = get_tweets_page(database_tweets, args.pagesize, num_page)
+    while len(tweets_page) != 0:
+        buffer_tweets = {}
+        for tweet in tweets_page:
             # Get output path and contents for the new CSV file
             csv_tweet_output_path =\
                 utils.create_tweet_output_path(tweet, output_dir)
@@ -43,10 +58,14 @@ for tweet in database_tweets.find():
             buffer_tweets[csv_tweet_output_path][0] += csv_tweet_contents
             buffer_tweets[csv_tweet_output_path][1] += 1
 
-# Perform the write operations in each of the files
-for output_path in buffer_tweets.keys():
+        # Perform the write operations in each of the files
+        for output_path in buffer_tweets.keys():
             with open(output_path, "a") as tweet_writer:
                 tweet_writer.write(buffer_tweets[output_path][0])
                 utils.increase_metadata_count(
                     os.path.join(output_dir, ".metadata.json"),
                     output_path, increase=buffer_tweets[output_path][1])
+
+        num_page += 1
+        tweets_page =\
+            get_tweets_page(database_tweets, args.pagesize, num_page)