process_tracker_python-83 Add extract file type and compression type

✨ Enabled file and compression types on ExtractTracker object

Author: Alex Meadows

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', ',', '"', '/');

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', ',', '"', '/');

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:

process_tracker/models/extract.py

@@ -119,17 +119,13 @@ class Extract(Base):
     extract_compression_type_id = Column(
         Integer,
         ForeignKey(
-            "process_tracker.extract_compression_type_lkup.extract_compression_type_id",
-            schema="process_tracker",
+            "process_tracker.extract_compression_type_lkup.extract_compression_type_id"
         ),
         nullable=True,
     )
     extract_filetype_id = Column(
         Integer,
-        ForeignKey(
-            "process_tracker.extract_filetype.extract_filetype_id",
-            schema="process_tracker",
-        ),
+        ForeignKey("process_tracker.extract_filetype_lkup.extract_filetype_id"),
         nullable=True,
     )
 

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))