fix: sensor to use scheduled_time & all the window params (#211)

* chore: cleanup pack-files step as it is not needed

* fix: sensor to use scheduled_time & all the window params

Author: sravankorumilli

.github/workflows/lint.yml

@@ -17,8 +17,6 @@ jobs:
       - uses: actions/checkout@v2
         with:
           fetch-depth: 0
-      - name: generate dependencies
-        run: make pack-files
       - name: golangci-lint
         uses: golangci/golangci-lint-action@v2
         with:

Makefile

@@ -7,22 +7,18 @@ LAST_TAG := "$(shell git rev-list --tags --max-count=1)"
 OPMS_VERSION := "$(shell git describe --tags ${LAST_TAG})-next"
 PROTON_COMMIT := "2a30976f7b40884ddd90e1792576c0941426e8bc"
 
-.PHONY: build test test-ci pack-files generate-proto unit-test-ci smoke-test integration-test vet coverage clean install lint
+.PHONY: build test test-ci generate-proto unit-test-ci smoke-test integration-test vet coverage clean install lint
 
 .DEFAULT_GOAL := build
 
-build: pack-files # build optimus binary
+build: # build optimus binary
 	@echo " > notice: skipped proto generation, use 'generate-proto' make command"
 	@echo " > building optimus version ${OPMS_VERSION}"
 	@go build -ldflags "-X ${NAME}/config.Version=${OPMS_VERSION} -X ${NAME}/config.BuildCommit=${LAST_COMMIT}" -o optimus .
 	@echo " - build complete"
 
 test-ci: smoke-test unit-test-ci vet ## run tests
 
-pack-files:
-	@echo " > packing resources"
-	@go generate ./..
-
 generate-proto: ## regenerate protos
 	@echo " > generating protobuf from odpf/proton"
 	@echo " > [info] make sure correct version of dependencies are installed using 'make install'"

ext/scheduler/airflow2/resources/__lib.py

@@ -174,8 +174,8 @@ def poke(self, context, session=None):
                                        '{} was deleted.'.format(self.upstream_dag))
 
         # calculate windows
-        execution_date = context['execution_date']
-        window_start, window_end = self.generate_window(execution_date, self.window_size)
+        schedule_time = context['next_execution_date']
+        window_start, window_end = self.generate_window(schedule_time)
         self.log.info(
             "consuming upstream window between: {} - {}".format(window_start.isoformat(), window_end.isoformat()))
         self.log.info("upstream interval: {}".format(dag_to_wait.schedule_interval))
@@ -208,12 +208,12 @@ def poke(self, context, session=None):
 
         return True
 
-    def generate_window(self, execution_date, window_size):
+    def generate_window(self, schedule_time):
         format_rfc3339 = "%Y-%m-%dT%H:%M:%SZ"
-        execution_date_str = execution_date.strftime(format_rfc3339)
+        schedule_time_str = schedule_time.strftime(format_rfc3339)
         # ignore offset & truncateto
-        task_window = JobSpecTaskWindow(window_size, 0, "m", self._optimus_client)
-        return task_window.get(execution_date_str)
+        task_window = JobSpecTaskWindow(self.window_size, self.window_offset, self.window_truncate_to, self._optimus_client)
+        return task_window.get(schedule_time_str)
 
     def _get_expected_upstream_executions(self, schedule_interval, window_start, window_end):
         expected_upstream_executions = []
@@ -325,24 +325,28 @@ def __init__(
             upstream_optimus_project: str,
             upstream_optimus_job: str,
             window_size: str,
+            window_offset: str,
+            window_truncate_to: str,
             *args,
             **kwargs) -> None:
         kwargs['mode'] = kwargs.get('mode', 'reschedule')
         super().__init__(**kwargs)
         self.optimus_project = upstream_optimus_project
         self.optimus_job = upstream_optimus_job
         self.window_size = window_size
+        self.window_offset = window_offset
+        self.window_truncate_to = window_truncate_to
         self._optimus_client = OptimusAPIClient(optimus_hostname)
 
     def poke(self, context):
-        execution_date = context['execution_date']
-        execution_date_str = execution_date.strftime(self.TIMESTAMP_FORMAT)
+        schedule_time = context['next_execution_date']
+        schedule_time_str = schedule_time.strftime(self.TIMESTAMP_FORMAT)
 
         # parse relevant metadata from the job metadata to build the task window
-        job_metadata = self._optimus_client.get_job_metadata(execution_date_str, self.optimus_project, self.optimus_job)
+        job_metadata = self._optimus_client.get_job_metadata(schedule_time_str, self.optimus_project, self.optimus_job)
         cron_schedule = lookup_non_standard_cron_expression(job_metadata['job']['interval'])
 
-        window_start, window_end = self.generate_window(execution_date, self.window_size)
+        window_start, window_end = self.generate_window(schedule_time)
         self.log.info(
             "consuming upstream window between: {} - {}".format(window_start.isoformat(), window_end.isoformat()))
         expected_upstream_executions = self._get_expected_upstream_executions(cron_schedule, window_start, window_end)
@@ -388,12 +392,11 @@ def _parse_datetime(self, timestamp) -> datetime:
         except ValueError:
             return datetime.strptime(timestamp, self.TIMESTAMP_MS_FORMAT)
 
-    def generate_window(self, execution_date, window_size):
+    def generate_window(self, schedule_time):
         format_rfc3339 = "%Y-%m-%dT%H:%M:%SZ"
-        execution_date_str = execution_date.strftime(format_rfc3339)
-        # ignore offset & truncate to
-        task_window = JobSpecTaskWindow(window_size, 0, "m", self._optimus_client)
-        return task_window.get(execution_date_str)
+        schedule_time_str = schedule_time.strftime(format_rfc3339)
+        task_window = JobSpecTaskWindow(self.window_size, self.window_offset, self.window_truncate_to, self._optimus_client)
+        return task_window.get(schedule_time_str)
 
 
 def optimus_failure_notify(context):

ext/scheduler/airflow2/resources/base_dag.py

@@ -195,7 +195,9 @@
     optimus_hostname="{{$.Hostname}}",
     upstream_optimus_project="{{$dependency.Project.Name}}",
     upstream_optimus_job="{{$dependency.Job.Name}}",
-    window_size="{{ $baseWindow.Size.String }}",
+    window_size={{$baseWindow.Size.String | quote}},
+    window_offset={{$baseWindow.Offset.String | quote}},
+    window_truncate_to={{$baseWindow.TruncateTo | quote}},
     poke_interval=SENSOR_DEFAULT_POKE_INTERVAL_IN_SECS,
     timeout=SENSOR_DEFAULT_TIMEOUT_IN_SECS,
     task_id="wait_{{$dependency.Job.Name | trunc 200}}-{{$dependencySchema.Name}}",

ext/scheduler/airflow2/resources/expected_compiled_template.py

@@ -198,6 +198,8 @@
     upstream_optimus_project="foo-external-project",
     upstream_optimus_job="foo-inter-dep-job",
     window_size="1h0m0s",
+    window_offset="0s",
+    window_truncate_to="d",
     poke_interval=SENSOR_DEFAULT_POKE_INTERVAL_IN_SECS,
     timeout=SENSOR_DEFAULT_TIMEOUT_IN_SECS,
     task_id="wait_foo-inter-dep-job-bq",