added couple of docstrings

added gcp storage intial code

Author: 116davinder

cloud/aws.py

@@ -4,11 +4,15 @@
 import os
 import time
 import threading
-from common import common
+from common import common, checkpoint
 
 class Upload:
 
     def s3_upload_file(s3_client,bucket,file_name,object_name):
+        """It will upload given `file_name` to given s3 `bucket`
+        and `object_name` s3 path.
+        """
+
         try:
             response = s3_client.upload_file(file_name,bucket,object_name)
             logging.info(f"upload successful at s3://{bucket}/{object_name}")
@@ -19,6 +23,12 @@ def s3_upload_file(s3_client,bucket,file_name,object_name):
             logging.error(f"{file_path} upload failed error {e}")
 
     def s3_upload(bucket,dir,topic_name,retry_upload_seconds,thread_count):
+        """It will initialize s3 client and 
+        based on checkpoint file for each partition.
+        It will call `s3_upload_file` function to upload.
+        
+        It will run after every `retry_upload_seconds`"""
+
         s3_client = boto3.client('s3')
         while True:
             _topic_dir = os.path.join(dir, topic_name)
@@ -43,27 +53,9 @@ def s3_upload(bucket,dir,topic_name,retry_upload_seconds,thread_count):
 
 class Download:
 
-    def s3_read_checkpoint_partition(dir,topic,partition):
-        try:
-            logging.debug(os.path.join(dir,topic,partition,"checkpoint"))
-            with open(os.path.join(dir,topic,partition,"checkpoint")) as c:
-                line = c.readline().strip()
-                _ck_file = line.split()[0]
-                _total_files = line.split()[1]
-                return {"checkpoint": _ck_file, "total_files": _total_files}
-        except FileNotFoundError as e:
-            logging.debug(e)
-        return None
-
-    def s3_write_checkpoint_partition(dir,topic,partition,msg):
-        try:
-            with open(os.path.join(dir,topic,partition,"checkpoint"), "w") as cp:
-                cp.write(msg)
-        except TypeError as e:
-            logging.error(e)
-            logging.error(f"writing checkpoint {msg} for {topic}, {partition} is failed")
-
     def s3_count_partitions(s3_client,bucket,topic):
+        """It will return number of objects in a given s3 bucket and s3 bucket path."""
+
         try:
             return s3_client.list_objects_v2(
                 Bucket=bucket,
@@ -75,6 +67,8 @@ def s3_count_partitions(s3_client,bucket,topic):
             exit(1)
 
     def s3_list_files(s3_client,bucket,path):
+        """It will list all files for given s3 bucket and s3 bucket path."""
+
         _list = []
         paginator = s3_client.get_paginator('list_objects_v2')
         operation_parameters = {'Bucket': bucket,'Prefix': path}
@@ -88,6 +82,21 @@ def s3_list_files(s3_client,bucket,path):
         return sorted(_list)
 
     def s3_download_file(s3_client,bucket,object_path,file_path):
+        """It will download two files .tar.gz and .tar.gz.sha256 .
+
+        Parameters
+        ----------
+        s3_client : boto3.client('s3')
+
+        bucket: str
+        
+        object_path: str
+            Description: path in s3 bucket
+
+        file_path: str
+            Description: path from local filesystem
+        """
+
         try:
             # donwload .tar.gz
             s3_client.download_file(bucket, object_path, file_path)
@@ -98,6 +107,13 @@ def s3_download_file(s3_client,bucket,object_path,file_path):
             logging.error(f"{file_path} failed with error {e}")
 
     def s3_download(bucket,topic,tmp_dir,retry_download_seconds=60):
+        """It will initialize s3 client and 
+        based on checkpoint file for each partition.
+        It will call `s3_download_file` function to download backup
+        and backup sha file.
+        
+        It will run after every `retry_download_seconds`"""
+
         s3_client = boto3.client('s3')
         while True:
             _pc = Download.s3_count_partitions(s3_client,bucket,topic)
@@ -106,7 +122,8 @@ def s3_download(bucket,topic,tmp_dir,retry_download_seconds=60):
                 os.makedirs(os.path.join(tmp_dir,topic,str(p)),exist_ok=True)
 
             for _pt in range(_pc):
-                _ck = Download.s3_read_checkpoint_partition(tmp_dir,topic,str(_pt))
+
+                _ck = checkpoint.read_checkpoint_partition(tmp_dir,topic,str(_pt)
                 _partition_path = os.path.join(topic,str(_pt))
                 _s3_partition_files = Download.s3_list_files(s3_client,bucket,_partition_path)
                 if _ck is not None:
@@ -134,7 +151,7 @@ def s3_download(bucket,topic,tmp_dir,retry_download_seconds=60):
                         Download.s3_download_file(s3_client,bucket,file,os.path.join(tmp_dir,file))
                         if file.endswith(".tar.gz"):
                             _ck['total_files'] += 1
-                            Download.s3_write_checkpoint_partition(tmp_dir,topic,str(_pt),file + " " + str(_ck['total_files']))
+                            checkpoint.write_checkpoint_partition(tmp_dir,topic,str(_pt),file + " " + str(_ck['total_files']))
 
             if _pc == 0:
                 logging.error(f"No Partitions found in given S3 path s3://{bucket}/{topic} retry seconds {retry_download_seconds}s")

cloud/gcp.py

@@ -0,0 +1,53 @@
+import logging
+import os
+import time
+import threading
+from common import common
+from google.cloud import storage
+
+class gUpload:
+
+    def gcs_upload_file(gcs_client,bucket,file_name,object_name):
+
+        blob = bucket.blob(object_name)
+        blob.upload_from_filename(file_name)
+        logging.info(f"upload successful at gcs://{bucket}/{object_name}")
+        if not file_name.endswith(".bin"):
+            logging.debug(f"deleting uploaded file {file_name}")
+            os.remove(file_name)
+
+    def gcs_upload(bucket,dir,topic_name,retry_upload_seconds,thread_count):
+        gcs_client = storage.client()
+        _bucket = gcs_client.get_bucket(bucket)
+        while True:
+            _topic_dir = os.path.join(dir, topic_name)
+            _count_partition_dirs = len(common.listDirs(_topic_dir))
+            _list = common.findFilesInFolder(_topic_dir)
+            if len(_list) > _count_partition_dirs and threading.active_count() <= thread_count:
+                for file_name in _list:
+                    file_name = str(file_name)
+                    file_size = os.path.getsize(file_name)
+                    if ( file_size > 0 and file_name.endswith(".tar.gz")
+                        or
+                        file_size > 0 and file_name.endswith(".tar.gz.sha256")):
+                        object_name = file_name.split(dir)[1]
+                        t = threading.Thread(
+                            target=gUpload.gcs_upload_file,
+                            args=[gcs_client,bucket,file_name,object_name],
+                            name="GCS Upload Threads"
+                        ).start()
+            else:
+                logging.info(f"gcs upload retry for new files in {retry_upload_seconds} seconds")
+                time.sleep(retry_upload_seconds)
+
+class gDownload:
+
+    def gcs_list_blobs_with_prefix(gcs_client,bucket,start_offset=None,end_offset=None,prefix=None):
+        return gcs_client.list_blobs(
+            bucket,
+            prefix=prefix,
+            delimiter="/",
+            start_offset=start_offset,
+            end_offset=end_offset
+        )
+        

common/checkpoint.py

@@ -0,0 +1,26 @@
+import os
+import logging
+
+def read_checkpoint_partition(dir,topic,partition):
+    """It will read checkpoint file for given topic and partition."""
+
+    try:
+        logging.debug(os.path.join(dir,topic,partition,"checkpoint"))
+        with open(os.path.join(dir,topic,partition,"checkpoint")) as c:
+            line = c.readline().strip()
+            _ck_file = line.split()[0]
+            _total_files = line.split()[1]
+            return {"checkpoint": _ck_file, "total_files": _total_files}
+    except FileNotFoundError as e:
+        logging.error(e)
+    return None
+
+def write_checkpoint_partition(dir,topic,partition,msg):
+    """It will write checkpoint message for given topic and partition to given checkpoint file"""
+
+    try:
+        with open(os.path.join(dir,topic,partition,"checkpoint"), "w") as cp:
+            cp.write(msg)
+    except TypeError as e:
+        logging.error(e)
+        logging.error(f"writing checkpoint {msg} for {topic}, {partition} is failed")

requirements.txt

@@ -1,2 +1,3 @@
 confluent_kafka
 boto3
+google-cloud-storage