process_tracker_python-7 Register extracts by location

OpenDataAlex · OpenDataAlex · commit b6ce2abfd4e2 · 2019-05-18T09:04:55.000-04:00
✨ Can now register extracts by given location

Have tested new extract registration by given location for local directory paths.
diff --git a/Pipfile b/Pipfile
@@ -5,11 +5,12 @@ verify_ssl = true
 
 [dev-packages]
 coverage="*"
+moto="*"
 python-coveralls="*"
 coveralls="*"
-twine="*"
 
 [packages]
+boto3="*"
 sqlalchemy="*"
 sqlalchemy-utils="*"
 python-dateutil="*"
diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/process_tracker/data_store.py b/process_tracker/data_store.py
@@ -59,35 +59,32 @@ def verify_and_connect_to_data_store(self):
         data_store_port = os.environ.get('process_tracking_data_store_port')
         data_store_name = os.environ.get('process_tracking_data_store_name')
 
-        data_store_error_flag = False
+        errors = []
 
         if data_store_type is None:
-            raise Exception('Data store type is not set.')
-            data_store_error_flag = True
+            errors.append(Exception('Data store type is not set.'))
 
         if data_store_username is None:
-            raise Exception('Data store username is not set.')
-            data_store_error_flag = True
+            errors.append(Exception('Data store username is not set.'))
 
         if data_store_password is None:
-            raise Exception('Data store password is not set')
-            data_store_error_flag = True
+            errors.append(Exception('Data store password is not set'))
 
         if data_store_host is None:
-            raise Exception('Data store host is not set')
-            data_store_error_flag = True
+            errors.append(Exception('Data store host is not set'))
 
         if data_store_port is None:
-            raise Exception('Data store port is not set')
-            data_store_error_flag = True
+            errors.append(Exception('Data store port is not set'))
 
         if data_store_name is None:
-            raise Exception('Data store name is not set')
-            data_store_error_flag = True
+            errors.append(Exception('Data store name is not set'))
 
-        if data_store_error_flag:
-            raise Exception('Data store has not been properly configured.  Please read how to set up the Process '
-                            'Tracking data store by going to: <insert read the docs url here>')
+        if errors:
+
+            errors.append(Exception('Data store has not been properly configured.  Please read how to set up the Process '
+                                    'Tracking data store by going to: <insert read the docs url here>'))
+
+            raise Exception(errors)
 
         relational_stores = ['postgresql']
         nonrelational_stores = []
diff --git a/process_tracker/extract_tracker.py b/process_tracker/extract_tracker.py
@@ -1,49 +1,61 @@
 # Extract Tracking
 # Used in the creation and editing of extract records.  Used in conjunction with process tracking.
 from datetime import datetime
-from os.path import basename, join, normpath
+import logging
+from os.path import join
 
-from process_tracker.data_store import DataStore
+from sqlalchemy.orm import Session
 
+from process_tracker.data_store import DataStore
+from process_tracker.location_tracker import LocationTracker
 from process_tracker.models.extract import Extract, ExtractProcess, ExtractStatus, Location
 
 
 class ExtractTracker:
 
-    def __init__(self, process_run, filename, location_path, location_name=None):
+    def __init__(self, process_run, filename, location=None, location_name=None, location_path=None, status=None):
         """
         ExtractTracker is the primary engine for tracking data extracts
         :param process_run: The process object working with extracts (either creating or consuming)
         :type process_run: ProcessTracker object
         :param filename: Name of the data extract file.
         :type filename:  string
+        :param location: SQLAlchemy Location object
         :param location_path: Location (filepath, s3 bucket, etc.) where the file is stored
         :type location_path: string
         :param location_name: Optional parameter to provide a specific name for the location.  If not provided, will use
                               the last directory in the path as the location name.  If type of location can be
                               determined (i.e. S3 bucket), the location type will be prepended.
         :type location_name: string
+        :param status: Optional if status does not need to be 'initializing', which is default.
+        :type status: string
         """
+        self.logger = logging.getLogger(__name__)
+
         self.data_store = DataStore()
         self.session = self.data_store.session
         self.process_run = process_run
 
-        if location_name is None:
-            location_name = self.derive_location_name(location_path=location_path)
-
         self.source = self.process_run.source
         self.filename = filename
-        self.full_filename = join(location_path, filename)
 
-        self.location = self.data_store.get_or_create(model=Location
-                                                      , location_name=location_name
-                                                      , location_path=location_path)
+        if location is not None:
+            self.location = location
+        elif location_path is not None:
+            self.location = LocationTracker(location_name=location_name, location_path=location_path)
+        else:
+            raise Exception('A location object or location_path must be provided.')
 
         self.extract = self.data_store.get_or_create(model=Extract
                                                      , extract_filename=filename
-                                                     , extract_location_id=self.location.location_id
+                                                     , extract_location_id=self.location.location.location_id
                                                      , extract_source_id=self.source.source_id)
 
+        if location_path is not None:
+            self.full_filename = join(location_path, filename)
+        else:
+            self.full_filename = join(self.location.location_path, self.extract.extract_filename)
+
         # Getting all status types in the event there are custom status types added later.
         self.extract_status_types = self.get_extract_status_types()
 
@@ -59,7 +71,16 @@ def __init__(self, process_run, filename, location_path, location_name=None):
 
         self.extract_process = self.retrieve_extract_process()
 
-        self.extract.extract_status_id = self.extract_status_initializing
+        if status is not None and self.extract_status_types[status]:
+            self.logger.info('Setting extract status to %s' % status)
+            self.extract.extract_status_id = self.extract_status_types[status]
+            self.extract_process.extract_process_status_id = self.extract_status_types[status]
+        else:
+            if status is not None:
+                self.logger.error('Provided status %s is not in extract_status_types_lkup.  '
+                                  'Setting to initializing.' % status)
+            self.extract.extract_status_id = self.extract_status_initializing
+
         self.session.commit()
 
     def change_extract_status(self, new_status):
@@ -82,34 +103,6 @@ def change_extract_status(self, new_status):
             raise Exception('%s is not a valid extract status type.  '
                             'Please add the status to extract_status_lkup' % new_status)
 
-    @staticmethod
-    def derive_location_name(location_path):
-        """
-        If location name is not provided, attempt to derive name from path.
-        :param location_path: The data extract file location path.
-        :return:
-        """
-        # Idea is to generalize things like grabbing the last directory name in the path,
-        # what type of path is it (normal, s3, etc.)
-
-        location_prefix = None
-        location_name = ""
-
-        location_path = location_path.lower()  # Don't care about casing.
-
-        if "s3" in location_path:
-            # If the path is an S3 Bucket, prefix to name.
-
-            location_prefix = "s3"
-
-        if location_prefix is not None:
-
-            location_name = location_prefix + " - "
-
-        location_name += basename(normpath(location_path))
-
-        return location_name
-
     def get_extract_status_types(self):
         """
         Get list of process status types and return dictionary.
diff --git a/process_tracker/location_tracker.py b/process_tracker/location_tracker.py
@@ -0,0 +1,68 @@
+# Location
+# For processes dealing with Extract Locations.
+
+from os.path import basename, normpath
+
+from process_tracker.data_store import DataStore
+from process_tracker.models.extract import Location, LocationType
+
+
+class LocationTracker:
+
+    def __init__(self, location_path, location_name=None):
+
+        self.data_store = DataStore()
+
+        self.location_path = location_path.lower()
+        self.location_type = self.derive_location_type()
+
+        if location_name is None:
+            self.location_name = self.derive_location_name()
+        else:
+            self.location_name = location_name
+
+        self.location = self.data_store.get_or_create(model=Location
+                                                      , location_name=self.location_name
+                                                      , location_path=location_path
+                                                      , location_type=self.location_type.location_type_id)
+
+    def derive_location_name(self):
+        """
+        If location name is not provided, attempt to derive name from path.
+        :return:
+        """
+        # Idea is to generalize things like grabbing the last directory name in the path,
+        # what type of path is it (normal, s3, etc.)
+
+        location_prefix = None
+
+        location_name = ""
+
+        if "s3" in self.location_path:
+            # If the path is an S3 Bucket, prefix to name.
+
+            location_prefix = "s3"
+
+        if location_prefix is not None:
+            location_name = location_prefix + " - "
+
+        location_name += basename(normpath(self.location_path))
+
+        return location_name
+
+    def derive_location_type(self):
+        """
+        Determine the type of location provided.
+        :return:
+        """
+
+        if "s3" in self.location_path:
+
+            location_type = self.data_store.get_or_create(model=LocationType
+                                                          , location_type_name="s3")
+
+        else:
+            location_type = self.data_store.get_or_create(model=LocationType
+                                                          , location_type_name="local directory")
+
+        return location_type
diff --git a/process_tracker/models/extract.py b/process_tracker/models/extract.py
@@ -50,12 +50,25 @@ class ExtractProcess(Base):
     extract_processes = relationship('ProcessTracking', foreign_keys=[process_tracking_id])
 
 
+class LocationType(Base):
+
+    __tablename__ = "location_type_lkup"
+
+    location_type_id = Column(Integer, Sequence('location_type_lkup_location_type_id_seq'), primary_key=True)
+    location_type_name = Column(String(25), unique=True, nullable=False)
+
+    locations = relationship('Location', back_populates='location_types')
+    
+
 class Location(Base):
 
     __tablename__ = "location_lkup"
 
     location_id = Column(Integer, Sequence('location_lkup_location_id_seq'), primary_key=True)
     location_name = Column(String(750), nullable=False, unique=True)
     location_path = Column(String(750), nullable=False, unique=True)
+    location_type = Column(Integer, ForeignKey("location_type_lkup.location_type_id"))
 
     extracts = relationship("Extract")
+
+    location_types = relationship('LocationType', foreign_keys=[location_type])
diff --git a/process_tracker/process_tracker.py b/process_tracker/process_tracker.py
@@ -3,9 +3,14 @@
 
 from datetime import datetime
 import logging
+import os
 from os.path import join
 
+import boto3
+
 from process_tracker.data_store import DataStore
+from process_tracker.extract_tracker import ExtractTracker
+from process_tracker.location_tracker import LocationTracker
 
 from process_tracker.models.actor import Actor
 from process_tracker.models.extract import Extract, ExtractProcess, ExtractStatus, Location
@@ -214,6 +219,31 @@ def raise_run_error(self, error_type_name, error_description=None, fail_run=Fals
             self.session.commit()
             raise Exception('Process halting.  An error triggered the process to fail.')
 
+    def register_extracts_by_location(self, location_path, location_name=None):
+        """
+        For a given location, find all files and attempt to register them.
+        :param location_name: Name of the location
+        :param location_path: Path of the location
+        :return:
+        """
+        location = LocationTracker(location_path=location_path, location_name=location_name)
+
+        if location.location_type == "s3":
+            s3 = boto3.resource('s3')
+            bucket = s3.Bucket(location.location_path)
+
+            for file in bucket.objects.all():
+                ExtractTracker(process_run=self
+                               , filename=file
+                               , location=location
+                               , status='ready')
+        else:
+            for file in os.listdir(location_path):
+                ExtractTracker(process_run=self
+                               , filename=file
+                               , location=location
+                               , status='ready')
+
     def register_new_process_run(self):
         """
         When a new process instance is starting, register the run in process tracking.
@@ -252,7 +282,6 @@ def register_new_process_run(self):
 
         else:
             raise Exception('The process %s is currently running.' % self.process_name)
-            exit()
 
     def set_process_run_low_high_dates(self, low_date=None, high_date=None):
         """
diff --git a/setup.py b/setup.py
@@ -15,6 +15,7 @@
     packages=setuptools.find_packages(),
     test_suite='tests.process_tracker_test_suite',
     install_requires=[
+        'boto3 >= 1.9.150',
         'sqlalchemy >= 1.3.3',
         'sqlalchemy-utils >= 0.33.11',
         'python-dateutil >= 2.8.0',
@@ -23,9 +24,9 @@
     extras_requires={
         'dev': [
             'coverage >= "4.0.3',
+            'moto >= 1.3.8'
             'python-coveralls >= 2.9.1',
             'coveralls >= 1.7.0',
-            'twine >= 1.13.0'
         ]
     },
     classifiers=[
diff --git a/tests/test_extract_tracker.py b/tests/test_extract_tracker.py
@@ -94,8 +94,8 @@ def test_derive_location_name_local_path(self):
         :return:
         """
         extract = ExtractTracker(process_run=self.process_run
-                                      , filename='test_extract_filename2.csv'
-                                      , location_path='/home/test/extract_dir2')
+                                 , filename='test_extract_filename2.csv'
+                                 , location_path='/home/test/extract_dir2')
 
         location = self.session.query(Location).filter(Location.location_id == extract.extract.extract_location_id)
 
diff --git a/tests/test_location_tracker.py b/tests/test_location_tracker.py
diff --git a/tests/test_process_tracker.py b/tests/test_process_tracker.py
