From c1fccd65e96d2c9a5a6481fef4ae74c6b5c3b7a2 Mon Sep 17 00:00:00 2001
From: Sree Vaddi <441385+sreev@users.noreply.github.com>
Date: Tue, 13 Oct 2020 12:51:18 -0700
Subject: [PATCH] add run_id; update client and client_wo; add test (#111)

* add run_id; update client and client_wo; add test

Signed-off-by: SreeV <441385+sreev@users.noreply.github.com>

* revert test wo client

Signed-off-by: SreeV <441385+sreev@users.noreply.github.com>
---
 marquez_client/client.py    | 7 +++++--
 marquez_client/client_wo.py | 6 +++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/marquez_client/client.py b/marquez_client/client.py
index 4bdef2c..4ef4489 100644
--- a/marquez_client/client.py
+++ b/marquez_client/client.py
@@ -192,8 +192,8 @@ def tag_dataset_field(self, namespace_name, dataset_name, field_name,
 
     # Job API
     def create_job(self, namespace_name, job_name, job_type, location=None,
-                   input_dataset=None,
-                   output_dataset=None, description=None, context=None):
+                   input_dataset=None, output_dataset=None, description=None,
+                   context=None, run_id=None):
         Utils.check_name_length(namespace_name, 'namespace_name')
         Utils.check_name_length(job_name, 'job_name')
         Utils.is_instance_of(job_type, JobType)
@@ -204,6 +204,9 @@ def create_job(self, namespace_name, job_name, job_type, location=None,
             'type': job_type.name
         }
 
+        if run_id:
+            payload['runId'] = run_id
+
         if context:
             payload['context'] = context
 
diff --git a/marquez_client/client_wo.py b/marquez_client/client_wo.py
index d0e04cf..6d646ea 100644
--- a/marquez_client/client_wo.py
+++ b/marquez_client/client_wo.py
@@ -115,7 +115,8 @@ def create_dataset(self, namespace_name, dataset_name, dataset_type,
     # Job API
     def create_job(self, namespace_name, job_name, job_type,
                    location=None, input_dataset=None,
-                   output_dataset=None, description=None, context=None):
+                   output_dataset=None, description=None, context=None,
+                   run_id=None):
         Utils.check_name_length(namespace_name, 'namespace_name')
         Utils.check_name_length(job_name, 'job_name')
         Utils.is_instance_of(job_type, JobType)
@@ -126,6 +127,9 @@ def create_job(self, namespace_name, job_name, job_type,
             'type': job_type.name
         }
 
+        if run_id:
+            payload['runId'] = run_id
+
         if context:
             payload['context'] = context