process_tracker_python-138 Add file size to extract_tracker

✨  File size and measure added to extract_tracker

Closes: #138

Author: Alex Meadows

dbscripts/mysql_process_tracker.sql

@@ -136,6 +136,24 @@ create table location_lkup
 		foreign key (location_type_id) references location_type_lkup (location_type_id)
 );
 
+create table process_tracker.filesize_type_lkup
+(
+	filesize_type_id serial not null,
+		constraint filesize_type_lkup_pk
+			primary key,
+	filesize_type_name varchar(75) not null,
+	filesize_type_code char(2) not null
+);
+
+create unique index filesize_type_lkup_filesize_type_code_uindex
+	on process_tracker.filesize_type_lkup (filesize_type_code);
+
+create unique index filesize_type_lkup_filesize_type_name_uindex
+	on process_tracker.filesize_type_lkup (filesize_type_name);
+
+create unique index filesize_type_lkup_udx03
+	on process_tracker.filesize_type_lkup (filesize_type_code, filesize_type_name);
+
 
 create table extract_tracking
 (
@@ -153,6 +171,8 @@ create table extract_tracking
 	extract_load_record_count int null comment 'The record count of the data set when loading the data file.',
 	extract_compression_type_id int null,
 	extract_filetype_id int null,
+	extract_filesize numeric null,
+	extract_filesize_type_id int null,
 	constraint extract_filename
 		unique (extract_filename),
 	constraint extract_tracking_fk03
@@ -162,7 +182,9 @@ create table extract_tracking
 	constraint extract_tracking_ibfk_1
 		foreign key (extract_location_id) references process_tracker.location_lkup (location_id),
 	constraint extract_tracking_ibfk_2
-		foreign key (extract_status_id) references process_tracker.extract_status_lkup (extract_status_id)
+		foreign key (extract_status_id) references process_tracker.extract_status_lkup (extract_status_id),
+	constraint extract_tracking_fk06
+        foreign key (extract_filesize_type) references process_tracker.filesize_type_lkup (filesize_type_id)
 );
 
 create index extract_location_id

dbscripts/mysql_process_tracker_defaults.sql

@@ -43,4 +43,9 @@ INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code,
 INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (5, 'gte', 'greater than or equal');
 INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (6, 'not', 'not equal');
 INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (7, 'lke', 'like');
-INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (8, 'in', 'in set');
+INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (8, 'in', 'in set');
+
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (1, 'kilobytes', 'KB');
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (2, 'megabytes', 'MB');
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (3, 'gigabytes', 'GB');
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (4, 'bytes', 'B ');

dbscripts/postgresql_process_tracker.sql

@@ -438,6 +438,26 @@ create unique index extract_compression_type_lkup_extract_compression_type_uinde
 	on process_tracker.extract_compression_type_lkup (extract_compression_type);
 
 
+create table process_tracker.filesize_type_lkup
+(
+	filesize_type_id serial not null
+		constraint filesize_type_lkup_pk
+			primary key,
+	filesize_type_name varchar(75) not null,
+	filesize_type_code char(2) not null
+);
+
+alter table process_tracker.filesize_type_lkup owner to pt_admin;
+
+create unique index filesize_type_lkup_filesize_type_code_uindex
+	on process_tracker.filesize_type_lkup (filesize_type_code);
+
+create unique index filesize_type_lkup_filesize_type_name_uindex
+	on process_tracker.filesize_type_lkup (filesize_type_name);
+
+create unique index filesize_type_lkup_udx03
+	on process_tracker.filesize_type_lkup (filesize_type_code, filesize_type_name);
+
 
 create table extract_tracking
 (
@@ -461,9 +481,13 @@ create table extract_tracking
 	extract_load_low_date_time timestamp,
 	extract_load_high_date_time timestamp,
 	extract_load_record_count integer,
-	extract_compression_type_id integer
+	extract_compression_type_id integer,
 		constraint extract_tracking_fk04
 			references extract_compression_type_lkup,
+	extract_filesize numeric,
+	extract_filesize_type_id integer
+		constraint extract_tracking_fk06
+			references process_tracker.filesize_type_lkup,
 	extract_filetype_id integer
 		constraint extract_tracking_fk05
 			references extract_filetype_lkup

dbscripts/postgresql_process_tracker_defaults.sql

@@ -42,4 +42,9 @@ INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code,
 INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (5, 'gte', 'greater than or equal');
 INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (6, 'not', 'not equal');
 INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (7, 'lke', 'like');
-INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (8, 'in', 'in set');
+INSERT INTO process_tracker.filter_type_lkup (filter_type_id, filter_type_code, filter_type_name) VALUES (8, 'in', 'in set');
+
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (1, 'kilobytes', 'KB');
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (2, 'megabytes', 'MB');
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (3, 'gigabytes', 'GB');
+INSERT INTO process_tracker.filesize_type_lkup (filesize_type_id, filesize_type_name, filesize_type_code) VALUES (4, 'bytes', 'B ');

process_tracker/extract_tracker.py

@@ -1,6 +1,7 @@
 # Extract Tracking
 # Used in the creation and editing of extract records.  Used in conjunction with process tracking.
 from datetime import datetime
+import itertools
 import logging
 import os
 from pathlib import Path
@@ -43,6 +44,7 @@ def __init__(
         filetype=None,
         config_location=None,
         extract_id=None,
+        file_size=None,
     ):
         """
         ExtractTracker is the primary engine for tracking data extracts
@@ -68,6 +70,8 @@ def __init__(
         :param extract_id: If trying to work with a specific extract that's in process, provide the id and it will be
         reconstructed.
         :type extract_id: int
+        :param file_size: The size of the file (i.e. 2.21MB)
+        :type file_size: str
         """
         log_level = SettingsManager(
             config_location=config_location
@@ -114,6 +118,18 @@ def __init__(
             self.extract_process = self.retrieve_extract_process()
             self.sources = self.extract.extract_sources
 
+            if (
+                self.extract.extract_filesize is not None
+                and self.extract.extract_filesize_type is not None
+            ):
+
+                self.file_size = "%s %s" % (
+                    self.extract.extract_filesize,
+                    self.extract.extract_filesize_type,
+                )
+            else:
+                self.file_size = None
+
         else:
             if filename is None:
                 error_msg = "Filename must be provided."
@@ -230,6 +246,14 @@ def __init__(
                 self.logger.info("Status was not provided.  Initializing.")
                 self.extract.extract_status_id = self.extract_status_initializing
 
+            if file_size is not None:
+                split_filesize = self.file_size_splitter(file_size=file_size)
+
+                self.extract.extract_filesize = split_filesize[0]
+                self.extract.extract_filesize_type = split_filesize[1]
+
+            self.file_size = file_size
+
             self.session.commit()
 
     def add_dependency(self, dependency_type, dependency):
@@ -382,6 +406,47 @@ def extract_dependency_check(self, extracts=None):
         else:
             return False
 
+    def file_size_splitter(self, file_size):
+        """
+        Take provided file size and split the amount from the measure.
+        :param file_size: The provided file size with measure (i.e. 2.2GB)
+        :type file_size: str
+        :return:
+        """
+        file_size_split = dict()
+
+        amount = "".join(itertools.takewhile(str.isdigit, file_size))
+        self.logger.debug("Amount is now: %s" % amount)
+        measure = "".join([i for i in file_size if not i.isdigit() and i != "."])
+        self.logger.debug("Measure is: %s" % measure)
+        amount = int(amount)
+        self.logger.debug("Amount is now: %s" % amount)
+        if (
+            measure == "bytes"
+            or measure.lower() == "b"
+            or measure == ""
+            or measure is None
+        ):
+            amount = amount / 1048576  # converting bytes to gb
+            measure = "GB"
+        elif measure.lower() == "mb":
+
+            measure = "MB"
+        elif measure.lower() == "gb":
+            measure = "GB"
+        else:
+            error_msg = (
+                "Unsupported measure detected. Please provide file size in bytes, MB, or GB.  "
+                "Measure provided was: %s" % measure
+            )
+            self.logger.error(error_msg)
+            raise Exception(error_msg)
+
+        file_size_split[0] = amount
+        file_size_split[1] = measure
+
+        return file_size_split
+
     def get_dataset_types(self):
         """
         Get list of dataset types associated to extract and return list.

process_tracker/models/extract.py

@@ -9,6 +9,7 @@
     DateTime,
     ForeignKey,
     Integer,
+    Numeric,
     Sequence,
     String,
     UniqueConstraint,
@@ -128,9 +129,16 @@ class Extract(Base):
         ForeignKey("process_tracker.extract_filetype_lkup.extract_filetype_id"),
         nullable=True,
     )
+    extract_filesize = Column(Numeric, nullable=True)
+    extract_filesize_type_id = Column(
+        Integer,
+        ForeignKey("process_tracker.filesize_type_lkup.filesize_type_id"),
+        nullable=True,
+    )
 
     compression_type = relationship("ExtractCompressionType")
     extract_filetype = relationship("ExtractFileType")
+    extract_filesize_type = relationship("FileSizeType")
     extract_dataset_types = relationship(
         "ExtractDatasetType",
         back_populates="dataset_type_extracts",
@@ -317,6 +325,30 @@ def __repr__(self):
         )
 
 
+class FileSizeType(Base):
+    __tablename__ = "filesize_type_lkup"
+    __table_args__ = {"schema": "process_tracker"}
+
+    filesize_type_id = Column(
+        Integer,
+        Sequence("filesize_type_lkup_filesize_type_id_seq", schema="process_tracker"),
+        primary_key=True,
+        nullable=False,
+    )
+    filesize_type_name = Column(String(75), nullable=False, unique=True)
+    filesize_type_code = Column(String(2), nullable=False, unique=True)
+
+    UniqueConstraint(filesize_type_code, filesize_type_name)
+
+    def __repr__(self):
+
+        return "<FilesizeType id=%s, code=%s, name=%s>" % (
+            self.filesize_type_id,
+            self.filesize_type_code,
+            self.filesize_type_name,
+        )
+
+
 class LocationType(Base):
 
     __tablename__ = "location_type_lkup"

tests/test_extract_tracker.py

@@ -754,3 +754,61 @@ def test_ensure_nulls_caught_on_instantiation(self):
             ExtractTracker(process_run=self.extract.process_run)
 
         return self.assertTrue("Filename must be provided." in str(context.exception))
+
+    def test_file_size_splitter_bytes(self):
+        """
+        Testing that if provided a file size in bytes, the file size in GB will be returned.
+        :return:
+        """
+        given_result = self.extract.file_size_splitter(file_size="1048576B")
+        given_result = {given_result[0], given_result[1]}
+        expected_result = {1, "GB"}
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_file_size_splitter_mb(self):
+        """
+        Testing that if provided a file size in MB, the file size will not be modified.
+        :return:
+        """
+        given_result = self.extract.file_size_splitter(file_size="1024MB")
+        given_result = {given_result[0], given_result[1]}
+        expected_result = {1024, "MB"}
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_file_size_splitter_gb(self):
+        """
+        Testing that if provided a file size in GB, the file size will not be modified.
+        :return:
+        """
+        given_result = self.extract.file_size_splitter(file_size="50GB")
+        given_result = {given_result[0], given_result[1]}
+        expected_result = {50, "GB"}
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_file_size_splitter_no_measure(self):
+        """
+        Testing that if provided a file size without a measure, the file size will be assumed to be bytes and returned
+        in GB.
+        :return:
+        """
+        given_result = self.extract.file_size_splitter(file_size="1048576")
+        given_result = {given_result[0], given_result[1]}
+        expected_result = {1, "GB"}
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_file_size_splitter_invalid_measure(self):
+        """
+        Testing that if provided a file size with an invalid measure, an exception will be thrown.
+        :return:
+        """
+        with self.assertRaises(Exception) as context:
+            self.extract.file_size_splitter(file_size="1024ZXB")
+
+        self.assertTrue(
+            "Unsupported measure detected. Please provide file size in bytes, MB, or GB.  Measure provided was: ZXB"
+            in str(context.exception)
+        )