add support for (a) bigint for execution_time_ms (b) rows processed for each execution-step

Author: ChintanRaval

README.md

@@ -29,6 +29,7 @@ $ python -m dpo [options] <db-connection-string> <command> [command-args]
     - `previous-execution-id`: identifier of an existing execution, ideally as returned by the `get-last-successful-execution` command.
   - `complete-step`: Marks the completion of an existing execution's step. Returns nothing unless there's an error.
     - `step-id`: a GUID identifier of an existing execution's step as returned by the `init-step` command.
+    - `rows-processed`: an optional numeric value to indicate the number of rows processed during this step. supports a postgresql BIGINT type value.
   - `complete-execution`: Marks the completion of an existing execution. Returns nothing unless there's an error.
     - `execution-id`: a GUID identifier of an existing execution as returned by the `init-execution` command.
 

dpo/DataPipelineOrchestrator.py

@@ -46,7 +46,7 @@ def __process_compare_step_models_command(self):
         ).execute()
 
     def __process_complete_step_command(self):
-        CompleteStepCommand(self.args.db_connection_string, self.args.step_id).execute()
+        CompleteStepCommand(self.args.db_connection_string, self.args.step_id, self.args.rows_processed).execute()
 
     def __process_complete_execution_command(self):
         CompleteExecutionCommand(self.args.db_connection_string, self.args.execution_id).execute()
@@ -137,6 +137,15 @@ def __get_arguments(self):
             'step_id',
             metavar='step-id',
             help='an execution\'s step id as received using \'init-step\' command')
+        complete_step_command_parser.add_argument(
+            '-r', '--rows-processed',
+            action='store',
+            const=None,
+            default=None,
+            type=int,
+            nargs='?',
+            help='An optional numeric value to indicate the number of rows processed during this step. '
+                 'Supports a PostgreSQL BIGINT type value.')
 
         complete_execution_command_parser = subparsers.add_parser(
             'complete-execution', help='completes the given execution.')

dpo/DataRepository.py

@@ -99,7 +99,7 @@ def save_execution_step_models(self, execution_step_id, model_checksums):
 
         return execution_step
 
-    def complete_execution_step(self, execution_step_id):
+    def complete_execution_step(self, execution_step_id, rows_processed):
         session = self.session_maker()
 
         execution_step = session.query(ExecutionStepEntity) \
@@ -110,6 +110,7 @@ def complete_execution_step(self, execution_step_id):
         execution_step.completed_on = self.get_current_db_datetime_with_timezone()
         execution_step.execution_time_ms \
             = (execution_step.completed_on - execution_step.started_on).total_seconds() * 1000
+        execution_step.rows_processed = rows_processed
 
         session.commit()
 

dpo/alembic/versions/c5c34dd0b8f2_refactor_schema_to_support_steps_and_.py

@@ -33,7 +33,7 @@ def upgrade():
                     sa.Column('started_on', sa.DateTime(timezone=True),
                               server_default=sa.text('now()'), nullable=False),
                     sa.Column('completed_on', sa.DateTime(timezone=True), nullable=True),
-                    sa.Column('execution_time_ms', sa.Integer(), nullable=True),
+                    sa.Column('execution_time_ms', sa.BigInteger(), nullable=True),
                     sa.PrimaryKeyConstraint('execution_id'),
                     schema='dpo'
                     )
@@ -49,7 +49,8 @@ def upgrade():
                     sa.Column('started_on', sa.DateTime(timezone=True), server_default=sa.text('now()'),
                               nullable=False),
                     sa.Column('completed_on', sa.DateTime(timezone=True), nullable=True),
-                    sa.Column('execution_time_ms', sa.Integer(), nullable=True),
+                    sa.Column('execution_time_ms', sa.BigInteger(), nullable=True),
+                    sa.Column('rows_processed', sa.BigInteger(), nullable=True),
                     sa.ForeignKeyConstraint(['execution_id'], ['dpo.execution.execution_id'], ),
                     sa.PrimaryKeyConstraint('execution_step_id'),
                     schema='dpo'
@@ -83,7 +84,7 @@ def upgrade():
         '''
         INSERT INTO dpo.execution_step (
             execution_step_id, execution_id, created_on, updated_on,
-            step_name, status, started_on, completed_on, execution_time_ms
+            step_name, status, started_on, completed_on, execution_time_ms, rows_processed
         )
         SELECT
             uuid_generate_v4() AS execution_step_id
@@ -95,6 +96,7 @@ def upgrade():
             , (select min(sm.created_on) from dpo.fe9eed6d812f_execution_model sm where sm.execution_id = e.id and sm.type = m.type) AS started_on
             , NULL AS completed_on
             , NULL AS execution_time_ms
+            , NULL AS rows_processed
         FROM  dpo.fe9eed6d812f_execution e
         JOIN  dpo.fe9eed6d812f_execution_model m on e.id = m.execution_id
         GROUP BY e.id, e.status, e.created_on, m.type

dpo/commands/CompleteStepCommand.py

@@ -2,10 +2,11 @@
 
 
 class CompleteStepCommand(BaseCommand):
-    def __init__(self, db_connection_string, step_id, logger=None):
+    def __init__(self, db_connection_string, step_id, rows_processed, logger=None):
         super().__init__(db_connection_string, logger)
         self._step_id = step_id
+        self._rows_processed = rows_processed
 
     def execute(self):
-        execution_step = self.repository.complete_execution_step(self._step_id)
+        execution_step = self.repository.complete_execution_step(self._step_id, self._rows_processed)
         self.logger.debug('Completed Execution Step: ' + str(execution_step))

dpo/entities/execution_entity.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, DateTime, Integer, String
+from sqlalchemy import Column, DateTime, BigInteger, String
 from sqlalchemy.sql import func
 from sqlalchemy.dialects.postgresql import UUID
 import uuid
@@ -44,7 +44,7 @@ class ExecutionEntity(Shared.BaseEntity):
                           nullable=True)
 
     execution_time_ms = Column('execution_time_ms',
-                               Integer,
+                               BigInteger,
                                nullable=True)
 
     def __str__(self):

dpo/entities/execution_step_entity.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, DateTime, Integer, String, ForeignKey
+from sqlalchemy import Column, DateTime, BigInteger, String, ForeignKey
 from sqlalchemy.sql import func
 from sqlalchemy.dialects.postgresql import UUID
 import uuid
@@ -57,9 +57,13 @@ class ExecutionStepEntity(Shared.BaseEntity):
                           nullable=True)
 
     execution_time_ms = Column('execution_time_ms',
-                               Integer,
+                               BigInteger,
                                nullable=True)
 
+    rows_processed = Column('rows_processed',
+                            BigInteger,
+                            nullable=True)
+
     def __str__(self):
         return f'execution_step_id={self.execution_step_id}, ' \
                f'created_on={self.created_on}, ' \
@@ -69,4 +73,5 @@ def __str__(self):
                f'status={self.status}, ' \
                f'started_on={self.started_on}, ' \
                f'completed_on={self.completed_on}, ' \
-               f'execution_time_ms={self.execution_time_ms}.'
+               f'execution_time_ms={self.execution_time_ms}, ' \
+               f'rows_processed={self.rows_processed}.'

tests/integration/test_integration.sh

@@ -118,8 +118,9 @@ CompareStepModels () {
 
 CompleteStep () {
     local stepId=$1
+    local rowsProcessed=$2
 
-    $dpo $dpo_conn_str complete-step $stepId
+    $dpo $dpo_conn_str complete-step $stepId --rows-processed $rowsProcessed
 }
 
 CompleteExecution () {
@@ -134,10 +135,12 @@ CompleteExecution () {
 ExecuteAndAssert () {
     # Arrange
     local iter_no=$1
-    local expected_lastSuccessfulExecId=$2
-    local expected_changedLoadModels=$3
-    local expected_changedTransformModels=$4
-    local __resultvar=$5
+    local loadRowsProcessed=$2
+    local transformRowsProcessed=$3
+    local expected_lastSuccessfulExecId=$4
+    local expected_changedLoadModels=$5
+    local expected_changedTransformModels=$6
+    local __resultvar=$7
 
     # Act
     local execId=$(InitialiseExecution)
@@ -153,13 +156,13 @@ ExecuteAndAssert () {
     echo "  iter$iter_no loadStepId                            = $loadStepId"
     local changedLoadModels=$(CompareStepModels $loadStepId $lastSuccessfulExecId)
     echo "  iter$iter_no changedLoadModels                     = '$changedLoadModels'"
-    CompleteStep $loadStepId
+    CompleteStep $loadStepId $loadRowsProcessed
 
     local transformStepId=$(InitialiseStep $execId TRANSFORM $transformModelDirectory)
     echo "  iter$iter_no transformStepId                       = $transformStepId"
     local changedTransformModels=$(CompareStepModels $transformStepId $lastSuccessfulExecId)
     echo "  iter$iter_no changedTransformModels                = '$changedTransformModels'"
-    CompleteStep $transformStepId
+    CompleteStep $transformStepId $transformRowsProcessed
 
     CompleteExecution $execId
 
@@ -209,7 +212,7 @@ echo "transform_model_3" > "$transformModelDirectory/$transform_model_3.sql"
 iter1_expected_lastSuccessfulExecId="" # pass in empty string to skip test since we don't know the past state of the pipeline
 iter1_expected_changedLoadModels="$load_model_1 $load_model_2"
 iter1_expected_changedTransformModels="$transform_model_1 $transform_model_2 $transform_model_3"
-ExecuteAndAssert "1" "$iter1_expected_lastSuccessfulExecId" "$iter1_expected_changedLoadModels" "$iter1_expected_changedTransformModels" iter1_execId
+ExecuteAndAssert "1" "2147483647" "" "$iter1_expected_lastSuccessfulExecId" "$iter1_expected_changedLoadModels" "$iter1_expected_changedTransformModels" iter1_execId
 
 ###############
 # Execution 2 #
@@ -232,4 +235,4 @@ echo "transform_model_4" > "$transformModelDirectory/$transform_model_4.sql"
 iter2_expected_lastSuccessfulExecId="$iter1_execId"
 iter2_expected_changedLoadModels=""
 iter2_expected_changedTransformModels="$transform_model_2 $transform_model_4"
-ExecuteAndAssert "2" "$iter2_expected_lastSuccessfulExecId" "$iter2_expected_changedLoadModels" "$iter2_expected_changedTransformModels" iter2_execId
+ExecuteAndAssert "2" "9223372036854775807" "" "$iter2_expected_lastSuccessfulExecId" "$iter2_expected_changedLoadModels" "$iter2_expected_changedTransformModels" iter2_execId