Sanketika-Obsrv/issue-tracker#153: lint and cleanup

Author: ravismula

alerts.yaml

@@ -0,0 +1,78 @@
+{
+    "type": "connector",
+    "metadata": {
+        "id": "object-store-connector",
+        "name": "Object Store Connector",
+        "version": "1.0.0",
+        "tenant": "multiple",
+        "category": "File",
+        "description": "The Object Store Connector is used to move data from any Object Store to the Obsrv platform",
+        "technology": "python",
+        "runtime": "spark",
+        "licence": "MIT",
+        "owner": "Sunbird",
+        "main_class": null,
+        "main_program": "object_store_connector"
+    },
+    "connectors": [
+        {
+            "id": "aws-s3-connector",
+            "name": "AWS S3",
+            "description": "The AWS S3 Connector is used to move data from any S3 Bucket to the Obsrv platform",
+            "icon": "https://upload.wikimedia.org/wikipedia/commons/b/bc/Amazon-S3-Logo.svg",
+            "config": {
+                "source": {
+                    "type": "s3",
+                    "bucket": "",
+                    "prefix": "",
+                    "prefix_format": "%y/%m/%d/%H",
+                    "credentials": {
+                        "access_key": "",
+                        "secret_key": "",
+                        "region": ""
+                    },
+                    "max_retries": "<int>"
+                }
+            }
+        },
+        {
+            "id": "azure-blob-connector",
+            "name": "Azure Blob Store",
+            "description": "The Azure Blob Store Connector is used to move data from any Azure Blob Container to the Obsrv platform",
+            "icon": "https://upload.wikimedia.org/wikipedia/commons/f/fa/Microsoft_Azure.svg",
+            "config": {
+                "source": {
+                    "type": "azure_blob",
+                    "container": "",
+                    "prefix": "",
+                    "prefix_format": "%y/%m/%d/%H",
+                    "credentials": {
+                        "account_name": "",
+                        "account_key": ""
+                    },
+                    "max_retries": "<int>"
+                }
+            }
+        },
+        {
+            "id": "gcs-connector",
+            "name": "Google Cloud Storage",
+            "description": "The GCS Connector is used to move data from any Google Bucket to the Obsrv platform",
+            "icon": "https://upload.wikimedia.org/wikipedia/commons/thumb/5/51/Google_Cloud_logo.svg/512px-Google_Cloud_logo.svg.png",
+            "config": {
+                "source": {
+                    "type": "gcs",
+                    "bucket": "",
+                    "prefix": "",
+                    "prefix_format": "%y/%m/%d/%H",
+                    "credentials": {
+                        "project_id": "",
+                        "client_email": "",
+                        "private_key": ""
+                    }
+                },
+                "max_retries": "<int>"
+            }
+        }
+    ]
+}

metadata.json

@@ -0,0 +1,4 @@
+metrics:
+  - name: num_api_calls
+    alias: Number of API Calls
+    description: The number of API calls made to the cloud provider

metrics.yaml

@@ -1,10 +1,12 @@
 import os
-from obsrv.connector.batch import SourceConnector
+
 from connector import ObjectStoreConnector
+from obsrv.connector.batch import SourceConnector
+
 # from obsrv.utils import Config
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     connector = ObjectStoreConnector()
-    config_file_path = os.path.join(os.path.dirname(__file__), 'config/config.yaml')
+    config_file_path = os.path.join(os.path.dirname(__file__), "config/config.yaml")
 
-    SourceConnector.process(connector=connector, config_file_path=config_file_path)
+    SourceConnector.process(connector=connector, config_file_path=config_file_path)

object_store_connector/__main__.py

@@ -1,26 +1,24 @@
 import datetime
-import time
 import json
+import time
 from typing import Any, Dict, Iterator
+
+from models.object_info import ObjectInfo
 from obsrv.common import ObsrvException
+from obsrv.connector import ConnectorContext, MetricsCollector
 from obsrv.connector.batch import ISourceConnector
-from obsrv.connector import ConnectorContext
-from obsrv.connector import MetricsCollector
-from obsrv.models import ErrorData, StatusCode, ExecutionState
+from obsrv.models import ErrorData, ExecutionState, StatusCode
 from obsrv.utils import LoggerController
-
-from pyspark.sql import SparkSession, DataFrame
+from provider.s3 import S3
 from pyspark.conf import SparkConf
+from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.functions import lit
-from pyspark.sql.types import *
-
-from provider.s3 import S3
-from models.object_info import ObjectInfo
 
 logger = LoggerController(__name__)
 
 MAX_RETRY_COUNT = 10
 
+
 class ObjectStoreConnector(ISourceConnector):
     def __init__(self):
         self.provider = None
@@ -30,21 +28,34 @@ def __init__(self):
         self.error_state = StatusCode.FAILED.value
         self.running_state = ExecutionState.RUNNING.value
         self.not_running_state = ExecutionState.NOT_RUNNING.value
-        self.queued_state = ExecutionState.QUEUED.value 
-
-    def process(self, sc: SparkSession, ctx: ConnectorContext, connector_config: Dict[Any, Any], metrics_collector: MetricsCollector) -> Iterator[DataFrame]:
-        if (ctx.state.get_state("status", default_value=self.not_running_state) == self.running_state):
+        self.queued_state = ExecutionState.QUEUED.value
+
+    def process(
+        self,
+        sc: SparkSession,
+        ctx: ConnectorContext,
+        connector_config: Dict[Any, Any],
+        metrics_collector: MetricsCollector,
+    ) -> Iterator[DataFrame]:
+        if (
+            ctx.state.get_state("status", default_value=self.not_running_state)
+            == self.running_state
+        ):
             logger.info("Connector is already running. Skipping processing.")
             return
 
         ctx.state.put_state("status", self.running_state)
         ctx.state.save_state()
-        self.max_retries = connector_config["source"]["max_retries"] if "max_retries" in connector_config["source"] else MAX_RETRY_COUNT
+        self.max_retries = (
+            connector_config["source"]["max_retries"]
+            if "max_retries" in connector_config["source"]
+            else MAX_RETRY_COUNT
+        )
         self._get_provider(connector_config)
         self._get_objects_to_process(ctx, metrics_collector)
         for res in self._process_objects(sc, ctx, metrics_collector):
             yield res
-        
+
         last_run_time = datetime.datetime.now()
         ctx.state.put_state("status", self.not_running_state)
         ctx.state.put_state("last_run_time", last_run_time)
@@ -54,67 +65,99 @@ def get_spark_conf(self, connector_config) -> SparkConf:
         self._get_provider(connector_config)
         if self.provider is not None:
             return self.provider.get_spark_config(connector_config)
-        
+
         return SparkConf()
 
     def _get_provider(self, connector_config: Dict[Any, Any]):
-        if connector_config["source"]["type"] == "s3":    
+        if connector_config["source"]["type"] == "s3":
             self.provider = S3(connector_config)
         else:
-            ObsrvException(ErrorData("INVALID_PROVIDER", "provider not supported: {}".format(connector_config["source"]["type"])))
-
-    def _get_objects_to_process(self, ctx: ConnectorContext, metrics_collector: MetricsCollector) -> None:
+            ObsrvException(
+                ErrorData(
+                    "INVALID_PROVIDER",
+                    "provider not supported: {}".format(
+                        connector_config["source"]["type"]
+                    ),
+                )
+            )
+
+    def _get_objects_to_process(
+        self, ctx: ConnectorContext, metrics_collector: MetricsCollector
+    ) -> None:
         objects = ctx.state.get_state("to_process", list())
         if ctx.building_block is not None and ctx.env is not None:
             self.dedupe_tag = "{}-{}".format(ctx.building_block, ctx.env)
         else:
-            raise ObsrvException(ErrorData("INVALID_CONTEXT", "building_block or env not found in context"))
-        
-        if not len(objects):    
-            num_files_discovered = ctx.stats.get_stat('num_files_discovered', 0)
+            raise ObsrvException(
+                ErrorData(
+                    "INVALID_CONTEXT", "building_block or env not found in context"
+                )
+            )
+
+        if not len(objects):
+            num_files_discovered = ctx.stats.get_stat("num_files_discovered", 0)
             objects = self.provider.fetch_objects(ctx, metrics_collector)
             objects = self._exclude_processed_objects(ctx, objects)
             metrics_collector.collect("new_objects_discovered", len(objects))
             ctx.state.put_state("to_process", objects)
             ctx.state.save_state()
             num_files_discovered += len(objects)
-            ctx.stats.put_stat("num_files_discovered", num_files_discovered)  
+            ctx.stats.put_stat("num_files_discovered", num_files_discovered)
             ctx.stats.save_stats()
 
         self.objects = objects
 
-    def _process_objects(self, sc: SparkSession, ctx: ConnectorContext, metrics_collector: MetricsCollector) -> Iterator[DataFrame]:
-        num_files_processed = ctx.stats.get_stat('num_files_processed', 0)
+    def _process_objects(
+        self,
+        sc: SparkSession,
+        ctx: ConnectorContext,
+        metrics_collector: MetricsCollector,
+    ) -> Iterator[DataFrame]:
+        num_files_processed = ctx.stats.get_stat("num_files_processed", 0)
         for i in range(0, len(self.objects)):
             obj = self.objects[i]
             obj["start_processing_time"] = time.time()
-            columns = StructType([])
-            df = self.provider.read_object(obj.get("location"), sc=sc, metrics_collector=metrics_collector, file_format=ctx.data_format)
+            df = self.provider.read_object(
+                obj.get("location"),
+                sc=sc,
+                metrics_collector=metrics_collector,
+                file_format=ctx.data_format,
+            )
 
             if df is None:
                 obj["num_of_retries"] += 1
                 if obj["num_of_retries"] < self.max_retries:
                     ctx.state.put_state("to_process", self.objects[i:])
                     ctx.state.save_state()
                 else:
-                    if not self.provider.update_tag(object=obj, tags=[{"key": self.dedupe_tag, "value": self.error_state}], metrics_collector=metrics_collector):
+                    if not self.provider.update_tag(
+                        object=obj,
+                        tags=[{"key": self.dedupe_tag, "value": self.error_state}],
+                        metrics_collector=metrics_collector,
+                    ):
                         break
                 return
             else:
                 df = self._append_custom_meta(sc, df, obj)
-                obj["download_time"] = time.time()-obj.get("start_processing_time")
-                if not self.provider.update_tag(object=obj, tags=[{"key": self.dedupe_tag, "value": self.success_state}], metrics_collector=metrics_collector):
+                obj["download_time"] = time.time() - obj.get("start_processing_time")
+                if not self.provider.update_tag(
+                    object=obj,
+                    tags=[{"key": self.dedupe_tag, "value": self.success_state}],
+                    metrics_collector=metrics_collector,
+                ):
                     break
-                ctx.state.put_state("to_process", self.objects[i+1:])
+                ctx.state.put_state("to_process", self.objects[i + 1 :])
                 ctx.state.save_state()
                 num_files_processed += 1
-                ctx.stats.put_stat("num_files_processed",num_files_processed)
+                ctx.stats.put_stat("num_files_processed", num_files_processed)
                 obj["end_processing_time"] = time.time()
                 yield df
-        
+
         ctx.stats.save_stats()
 
-    def _append_custom_meta(self, sc: SparkSession, df: DataFrame, object: ObjectInfo) -> DataFrame:
+    def _append_custom_meta(
+        self, sc: SparkSession, df: DataFrame, object: ObjectInfo
+    ) -> DataFrame:
         addn_meta = {
             "location": object.get("location"),
             "file_size_kb": object.get("file_size_kb"),
@@ -123,7 +166,7 @@ def _append_custom_meta(self, sc: SparkSession, df: DataFrame, object: ObjectInf
             "end_processing_time": object.get("end_processing_time"),
             "file_hash": object.get("file_hash"),
             "num_of_retries": object.get("num_of_retries"),
-            "in_time": object.get("in_time")
+            "in_time": object.get("in_time"),
         }
         df = df.withColumn("_addn_source_meta", lit(json.dumps(addn_meta, default=str)))
         return df
@@ -134,4 +177,4 @@ def _exclude_processed_objects(self, ctx: ConnectorContext, objects):
             if not any(tag["key"] == self.dedupe_tag for tag in obj.get("tags")):
                 to_be_processed.append(obj)
 
-        return to_be_processed
+        return to_be_processed

object_store_connector/connector.py

@@ -1,24 +1,20 @@
 from dataclasses import dataclass, field
+from datetime import datetime
 from typing import List
 from uuid import uuid4
-from datetime import datetime
+
 
 @dataclass
 class Tag:
     key: str
     value: str
 
     def to_dict(self):
-        return {
-            'key': self.key,
-            'value': self.value
-        }
+        return {"key": self.key, "value": self.value}
 
     def to_aws(self):
-        return {
-            'Key': self.key,
-            'Value': self.value
-        }
+        return {"Key": self.key, "Value": self.value}
+
 
 @dataclass
 class ObjectInfo:
@@ -38,17 +34,17 @@ class ObjectInfo:
 
     def to_json(self):
         return {
-            'id': self.id,
-            'connector_id': self.connector_id,
-            'dataset_id': self.dataset_id,
-            'location': self.location,
-            'format': self.format,
-            'file_size_kb': self.file_size_kb,
-            'in_time': self.in_time,
-            'download_time': self.download_time,
-            'start_processing_time': self.start_processing_time,
-            'end_processing_time': self.end_processing_time,
-            'file_hash': self.file_hash,
-            'num_of_retries': self.num_of_retries,
-            'tags': [tag.__dict__ for tag in self.tags]
-        }
+            "id": self.id,
+            "connector_id": self.connector_id,
+            "dataset_id": self.dataset_id,
+            "location": self.location,
+            "format": self.format,
+            "file_size_kb": self.file_size_kb,
+            "in_time": self.in_time,
+            "download_time": self.download_time,
+            "start_processing_time": self.start_processing_time,
+            "end_processing_time": self.end_processing_time,
+            "file_hash": self.file_hash,
+            "num_of_retries": self.num_of_retries,
+            "tags": [tag.__dict__ for tag in self.tags],
+        }