fix: provide a mechanism to contact optimus server and upstream optimus server (#624)

Author: sravankorumilli

ext/scheduler/airflow2/resources/__lib.py

@@ -240,6 +240,7 @@ class SuperExternalTaskSensor(BaseSensorOperator):
     def __init__(
             self,
             optimus_hostname: str,
+            upstream_optimus_hostname: str,
             upstream_optimus_project: str,
             upstream_optimus_namespace: str,
             upstream_optimus_job: str,
@@ -255,6 +256,7 @@ def __init__(
         self.window_size = window_size
         self.window_version = window_version
         self._optimus_client = OptimusAPIClient(optimus_hostname)
+        self._upstream_optimus_client = OptimusAPIClient(upstream_optimus_hostname)
 
     def poke(self, context):
         log_start_event(context, EVENT_NAMES.get("SENSOR_START_EVENT"))
@@ -297,15 +299,15 @@ def get_last_upstream_times(self, schedule_time_of_current_job, upstream_schedul
 
     def get_schedule_interval(self, schedule_time):
         schedule_time_str = schedule_time.strftime(TIMESTAMP_FORMAT)
-        job_metadata = self._optimus_client.get_job_metadata(schedule_time_str, self.optimus_namespace, self.optimus_project, self.optimus_job)
+        job_metadata = self._upstream_optimus_client.get_job_metadata(schedule_time_str, self.optimus_namespace, self.optimus_project, self.optimus_job)
         upstream_schedule = lookup_non_standard_cron_expression(job_metadata['spec']['interval'])
         return upstream_schedule
 
     # TODO the api will be updated with getJobRuns even though the field here refers to scheduledAt
     #  it points to execution_date
     def _are_all_job_runs_successful(self, schedule_time_window_start, schedule_time_window_end) -> bool:
         try:
-            api_response = self._optimus_client.get_job_run(self.optimus_project, self.optimus_job, schedule_time_window_start, schedule_time_window_end)
+            api_response = self._upstream_optimus_client.get_job_run(self.optimus_project, self.optimus_job, schedule_time_window_start, schedule_time_window_end)
             self.log.info("job_run api response :: {}".format(api_response))
         except Exception as e:
             self.log.warning("error while fetching job runs :: {}".format(e))

ext/scheduler/airflow2/resources/base_dag.py

@@ -216,6 +216,7 @@
 {{- $dependencySchema := $dependency.Job.Task.Unit.Info }}
 wait_{{$dependency.Job.Name | replace "-" "__dash__" | replace "." "__dot__"}} = SuperExternalTaskSensor(
     optimus_hostname="{{$.Hostname}}",
+    upstream_optimus_hostname="{{$.Hostname}}",
     upstream_optimus_project="{{$dependency.Project.Name}}",
     upstream_optimus_namespace="{{$dependency.Job.NamespaceSpec.Name}}",
     upstream_optimus_job="{{$dependency.Job.Name}}",
@@ -231,7 +232,8 @@
 {{- range $_, $dependency := $.Job.ExternalDependencies.OptimusDependencies}}
 {{ $identity := print $dependency.Name "-" $dependency.ProjectName "-" $dependency.JobName }}
 wait_{{ $identity | replace "-" "__dash__" | replace "." "__dot__"}} = SuperExternalTaskSensor(
-    optimus_hostname="{{$dependency.Host}}",
+    optimus_hostname="{{$.Hostname}}",
+    upstream_optimus_hostname="{{$dependency.Host}}",
     upstream_optimus_project="{{$dependency.ProjectName}}",
     upstream_optimus_namespace="{{$dependency.NamespaceName}}",
     upstream_optimus_job="{{$dependency.JobName}}",

ext/scheduler/airflow2/resources/expected_compiled_template.py

@@ -215,6 +215,7 @@
 
 wait_foo__dash__intra__dash__dep__dash__job = SuperExternalTaskSensor(
     optimus_hostname="http://airflow.example.io",
+    upstream_optimus_hostname="http://airflow.example.io",
     upstream_optimus_project="foo-project",
     upstream_optimus_namespace="bar-namespace",
     upstream_optimus_job="foo-intra-dep-job",
@@ -227,6 +228,7 @@
 )
 wait_foo__dash__inter__dash__dep__dash__job = SuperExternalTaskSensor(
     optimus_hostname="http://airflow.example.io",
+    upstream_optimus_hostname="http://airflow.example.io",
     upstream_optimus_project="foo-external-project",
     upstream_optimus_namespace="bar-namespace",
     upstream_optimus_job="foo-inter-dep-job",
@@ -239,7 +241,8 @@
 )
 
 wait_external__dash__optimus__dash__foo__dash__external__dash__optimus__dash__project__dash__foo__dash__external__dash__optimus__dash__dep__dash__job = SuperExternalTaskSensor(
-    optimus_hostname="http://optimus.external.io",
+    optimus_hostname="http://airflow.example.io",
+    upstream_optimus_hostname="http://optimus.external.io",
     upstream_optimus_project="foo-external-optimus-project",
     upstream_optimus_namespace="bar-external-optimus-namespace",
     upstream_optimus_job="foo-external-optimus-dep-job",