Merge pull request #96 from OpenDataAlex/process_tracker_python-83

OpenDataAlex · web-flow · commit 28bdbf4a9fdb · 2019-09-30T13:25:15.000-04:00
Process tracker python 83
diff --git a/dbscripts/mysql_process_tracker.sql b/dbscripts/mysql_process_tracker.sql
diff --git a/dbscripts/mysql_process_tracker_defaults.sql b/dbscripts/mysql_process_tracker_defaults.sql
@@ -20,4 +20,8 @@ INSERT INTO process_tracker.error_type_lkup (error_type_id, error_type_name) VAL
 INSERT INTO process_tracker.process_type_lkup (process_type_id, process_type_name) VALUES (default, 'Extract');
 INSERT INTO process_tracker.process_type_lkup (process_type_id, process_type_name) VALUES (default, 'Load');
 
-INSERT INTO process_tracker.system_lkup (system_id, system_key, system_value) VALUES (default, 'version', '0.6.0');
+INSERT INTO process_tracker.system_lkup (system_id, system_key, system_value) VALUES (default, 'version', '0.7.0');
+
+INSERT INTO process_tracker.extract_compression_type_lkup (extract_compression_type_id, extract_compression_type) VALUES (1, 'zip');
+
+INSERT INTO process_tracker.extract_filetype_lkup (extract_filetype_id, extract_filetype_code, extract_filetype, delimiter_char, quote_char, escape_char) VALUES (1, 'csv', 'Comma Separated Values', ',', '"', '/');
diff --git a/dbscripts/postgresql_process_tracker.sql b/dbscripts/postgresql_process_tracker.sql
diff --git a/dbscripts/postgresql_process_tracker_defaults.sql b/dbscripts/postgresql_process_tracker_defaults.sql
@@ -21,4 +21,8 @@ INSERT INTO process_tracker.error_type_lkup (error_type_id, error_type_name) VAL
 INSERT INTO process_tracker.process_type_lkup (process_type_id, process_type_name) VALUES (default, 'Extract');
 INSERT INTO process_tracker.process_type_lkup (process_type_id, process_type_name) VALUES (default, 'Load');
 
-INSERT INTO process_tracker.system_lkup (system_id, system_key, system_value) VALUES (default, 'version', '0.6.0');
+INSERT INTO process_tracker.system_lkup (system_id, system_key, system_value) VALUES (default, 'version', '0.7.0');
+
+INSERT INTO process_tracker.extract_compression_type_lkup (extract_compression_type_id, extract_compression_type) VALUES (1, 'zip');
+
+INSERT INTO process_tracker.extract_filetype_lkup (extract_filetype_id, extract_filetype_code, extract_filetype, delimiter_char, quote_char, escape_char) VALUES (1, 'csv', 'Comma Separated Values', ',', '"', '/');
diff --git a/process_tracker/extract_tracker.py b/process_tracker/extract_tracker.py
@@ -2,6 +2,7 @@
 # Used in the creation and editing of extract records.  Used in conjunction with process tracking.
 from datetime import datetime
 import logging
+import os
 from pathlib import Path
 
 from sqlalchemy.orm import aliased
@@ -11,6 +12,8 @@
 from process_tracker.utilities import utilities
 from process_tracker.models.extract import (
     Extract,
+    ExtractCompressionType,
+    ExtractFileType,
     ExtractDatasetType,
     ExtractDependency,
     ExtractProcess,
@@ -27,6 +30,8 @@ def __init__(
         location_name=None,
         location_path=None,
         status=None,
+        compression_type=None,
+        filetype=None,
         config_location=None,
     ):
         """
@@ -44,6 +49,10 @@ def __init__(
         :type location_name: string
         :param status: Optional if status does not need to be 'initializing', which is default.
         :type status: string
+        :param compression_type: Optional compression format of the extract.
+        :type compression_type: String
+        :param filetype: Optional file type of the extract. Will try to be derived from the filename if not provided.
+        :type filetype: string
         :param config_location: Optional location for the process_tracker configuration file.
         :type config_location: string
         """
@@ -74,12 +83,54 @@ def __init__(
         else:
             raise Exception("A location object or location_path must be provided.")
 
+        if compression_type is not None:
+            self.logger.info("Finding compression type.")
+            try:
+                self.compression_type = self.data_store.get_or_create_item(
+                    model=ExtractCompressionType,
+                    create=False,
+                    extract_compression_type=compression_type,
+                )
+            except Exception:
+                error_msg = "%s is not a valid compression type." % compression_type
+                self.logger.error(error_msg)
+                raise Exception(error_msg)
+
+            self.compression_type_id = self.compression_type.extract_compression_type_id
+        else:
+            self.compression_type_id = None
+
+        if filetype is not None:
+            self.logger.info("File type provided.  Verifying it is a valid filetype.")
+            try:
+                self.filetype = self.data_store.get_or_create_item(
+                    model=ExtractFileType, create=False, extract_filetype=filetype
+                )
+            except Exception:
+                error_msg = "%s is not a valid file type." % filetype
+                self.logger.error(error_msg)
+                raise Exception(error_msg)
+        else:
+            # Need to try to determine the filetype based on the extension of the filename.
+            file_extension = os.path.splitext(filename)[1]
+            file_extension = file_extension.replace(".", "")
+            self.logger.info(
+                "Trying to find record for file extension: %s" % file_extension
+            )
+            self.filetype = self.data_store.get_or_create_item(
+                model=ExtractFileType,
+                create=False,
+                extract_filetype_code=file_extension,
+            )
+
         self.logger.info("Registering extract.")
 
         self.extract = self.data_store.get_or_create_item(
             model=Extract,
             extract_filename=filename,
             extract_location_id=self.location.location.location_id,
+            extract_compression_type_id=self.compression_type_id,
+            extract_filetype_id=self.filetype.extract_filetype_id,
         )
 
         if location_path is not None:
diff --git a/process_tracker/location_tracker.py b/process_tracker/location_tracker.py
@@ -2,7 +2,6 @@
 # For processes dealing with Extract Locations.
 import logging
 from pathlib import PurePath
-from os.path import basename, dirname, join, isdir, normpath
 
 from process_tracker.utilities.aws_utilities import AwsUtilities
 from process_tracker.utilities.logging import console
diff --git a/process_tracker/models/extract.py b/process_tracker/models/extract.py
@@ -41,6 +41,52 @@ def __repr__(self):
         )
 
 
+class ExtractCompressionType(Base):
+
+    __tablename__ = "extract_compression_type_lkup"
+    __table_args__ = {"schema": "process_tracker"}
+
+    extract_compression_type_id = Column(
+        Integer,
+        Sequence(
+            "extract_compression_type_extract_compression_type_id",
+            schema="process_tracker",
+        ),
+        primary_key=True,
+        nullable=False,
+    )
+    extract_compression_type = Column(String(25), unique=True, nullable=False)
+
+    def __repr__(self):
+
+        return "<Extract Compression Type name=%s>" % (self.extract_compression_type)
+
+
+class ExtractFileType(Base):
+
+    __tablename__ = "extract_filetype_lkup"
+    __table_args__ = {"schema": "process_tracker"}
+
+    extract_filetype_id = Column(
+        Integer,
+        Sequence("extract_filetype_extract_filetype_id_seq", schema="process_tracker"),
+        primary_key=True,
+        nullable=False,
+    )
+    extract_filetype_code = Column(String(5), nullable=False)
+    extract_filetype = Column(String(75), nullable=False, unique=True)
+    delimiter_char = Column(String(1), nullable=True)
+    quote_char = Column(String(1), nullable=True)
+    escape_char = Column(String(1), nullable=True)
+
+    def __repr__(self):
+
+        return "<Extract Filetype code=%s, name=%s>" % (
+            self.extract_filetype_code,
+            self.extract_filetype,
+        )
+
+
 class Extract(Base):
 
     __tablename__ = "extract_tracking"
@@ -70,6 +116,18 @@ class Extract(Base):
     extract_load_low_date_time = Column(DateTime, nullable=True)
     extract_load_high_date_time = Column(DateTime, nullable=True)
     extract_load_record_count = Column(Integer, nullable=True)
+    extract_compression_type_id = Column(
+        Integer,
+        ForeignKey(
+            "process_tracker.extract_compression_type_lkup.extract_compression_type_id"
+        ),
+        nullable=True,
+    )
+    extract_filetype_id = Column(
+        Integer,
+        ForeignKey("process_tracker.extract_filetype_lkup.extract_filetype_id"),
+        nullable=True,
+    )
 
     extract_dataset_types = relationship(
         "ExtractDatasetType",
diff --git a/tests/test_extract_tracker.py b/tests/test_extract_tracker.py
@@ -547,3 +547,62 @@ def test_set_extract_record_count_invalid_type(self):
         return self.assertTrue(
             "blarg is not a valid audit_type." in str(context.exception)
         )
+
+    def test_set_compression_type(self):
+        """Testing that when an extract is created with a compression type, the type is associated correctly."""
+
+        extract = ExtractTracker(
+            process_run=self.process_run,
+            filename="Compression Type File.csv",
+            location_name="Test Location",
+            location_path="/home/test/extract_dir",
+            compression_type="zip",
+        )
+
+        given_result = extract.compression_type.extract_compression_type
+        expected_result = "zip"
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_set_compression_type_invalid(self):
+        """Testing that when an extract is created with an invalid compression type, an error is thrown."""
+
+        with self.assertRaises(Exception) as context:
+            ExtractTracker(
+                process_run=self.process_run,
+                filename="Compression Type File.csv",
+                location_name="Test Location",
+                location_path="/home/test/extract_dir",
+                compression_type="zap",
+            )
+
+        self.assertTrue("zap is not a valid compression type" in str(context.exception))
+
+    def test_set_file_type(self):
+        """Testing that when an extract is create with a valid file type, the type is associated correctly."""
+        extract = ExtractTracker(
+            process_run=self.process_run,
+            filename="Compression Type File.csv",
+            location_name="Test Location",
+            location_path="/home/test/extract_dir",
+            filetype="Comma Separated Values",
+        )
+
+        given_result = extract.filetype.extract_filetype_code
+        expected_result = "csv"
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_set_file_type_invalid(self):
+        """Testing that when an extract is created with an invalid file type, an error is thrown."""
+
+        with self.assertRaises(Exception) as context:
+            ExtractTracker(
+                process_run=self.process_run,
+                filename="Compression Type File.csv",
+                location_name="Test Location",
+                location_path="/home/test/extract_dir",
+                filetype="zap",
+            )
+
+        self.assertTrue("zap is not a valid file type" in str(context.exception))
