enable point-in-time column lineage in python client (#2276)

Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>

Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>
Co-authored-by: Willy Lulciuc <willy@datakin.com>

Author: pawel-big-lebowski

clients/python/marquez_client/client.py

@@ -343,8 +343,18 @@ def get_job_run(self, run_id):
         Utils.is_valid_uuid(run_id, 'run_id')
         return self._get(self._url('/jobs/runs/{0}', run_id))
 
-    def get_column_lineage_by_dataset(self, namespace, dataset, depth=None, with_downstream=None):
-        node_id = "dataset:{0}:{1}".format(namespace, dataset)
+    def get_column_lineage_by_dataset(
+            self,
+            namespace,
+            dataset,
+            depth=None,
+            with_downstream=None,
+            version=None
+    ):
+        node_id = self._append_version_to_node_id(
+            "dataset:{0}:{1}".format(namespace, dataset),
+            version
+        )
         return self._get_column_lineage(node_id, depth, with_downstream)
 
     def get_column_lineage_by_dataset_field(
@@ -353,15 +363,29 @@ def get_column_lineage_by_dataset_field(
             dataset,
             field,
             depth=None,
-            with_downstream=None
+            with_downstream=None,
+            version=None,
     ):
-        node_id = "datasetField:{0}:{1}:{2}".format(namespace, dataset, field)
+        node_id = self._append_version_to_node_id(
+            "datasetField:{0}:{1}:{2}".format(namespace, dataset, field),
+            version
+        )
         return self._get_column_lineage(node_id, depth, with_downstream)
 
-    def get_column_lineage_by_job(self, namespace, job, depth=None, with_downstream=None):
-        node_id = "job:{0}:{1}".format(namespace, job)
+    def get_column_lineage_by_job(self, namespace, job, depth=None,
+                                  with_downstream=None, version=None):
+        node_id = self._append_version_to_node_id(
+            "job:{0}:{1}".format(namespace, job),
+            version
+        )
         return self._get_column_lineage(node_id, depth, with_downstream)
 
+    def _append_version_to_node_id(self, node_id, version):
+        if version is not None:
+            return node_id + "#" + version
+        else:
+            return node_id
+
     def _get_column_lineage(self, node_id, depth, with_downstream):
         return self._get(
             self._url('/column-lineage'),

clients/python/tests/test_marquez_client.py

@@ -979,21 +979,16 @@ def test_get_column_lineage_by_dataset(mock_get, client):
     mock_get.return_value.json.return_value = COLUMN_LINEAGE
 
     column_lineage = client.get_column_lineage_by_dataset(
-        "namespace_a",
-        "dataset_a",
-        DEFAULT_DEPTH,
-        DEFAULT_WITH_DOWNSTREAM
+        namespace="namespace_a", dataset="dataset_a", depth=DEFAULT_DEPTH,
+        with_downstream=DEFAULT_WITH_DOWNSTREAM, version="some-version"
     )
 
     assert column_lineage == COLUMN_LINEAGE
-
     mock_get.assert_called_once_with(
-        url=client._url(
-            '/column-lineage'
-        ),
+        url=client._url('/column-lineage'),
         headers=mock.ANY,
         params={
-            'nodeId': 'dataset:namespace_a:dataset_a',
+            'nodeId': 'dataset:namespace_a:dataset_a#some-version',
             'depth': DEFAULT_DEPTH,
             'withDownstream': DEFAULT_WITH_DOWNSTREAM
         },
@@ -1005,24 +1000,17 @@ def test_get_column_lineage_by_dataset(mock_get, client):
 def test_get_column_lineage_by_dataset_field(mock_get, client):
     mock_get.return_value.status_code.return_value = HTTPStatus.OK
     mock_get.return_value.json.return_value = COLUMN_LINEAGE
-
     column_lineage = client.get_column_lineage_by_dataset_field(
-        "namespace_a",
-        "dataset_a",
-        "field_a",
-        DEFAULT_DEPTH,
-        DEFAULT_WITH_DOWNSTREAM
+        namespace="namespace_a", dataset="dataset_a", field="field_a", depth=DEFAULT_DEPTH,
+        with_downstream=DEFAULT_WITH_DOWNSTREAM, version="some-version"
     )
 
     assert column_lineage == COLUMN_LINEAGE
-
     mock_get.assert_called_once_with(
-        url=client._url(
-            '/column-lineage'
-        ),
+        url=client._url('/column-lineage'),
         headers=mock.ANY,
         params={
-            'nodeId': 'datasetField:namespace_a:dataset_a:field_a',
+            'nodeId': 'datasetField:namespace_a:dataset_a:field_a#some-version',
             'depth': DEFAULT_DEPTH,
             'withDownstream': DEFAULT_WITH_DOWNSTREAM
         },
@@ -1034,23 +1022,17 @@ def test_get_column_lineage_by_dataset_field(mock_get, client):
 def test_get_column_lineage_by_job(mock_get, client):
     mock_get.return_value.status_code.return_value = HTTPStatus.OK
     mock_get.return_value.json.return_value = COLUMN_LINEAGE
-
     column_lineage = client.get_column_lineage_by_job(
-        "namespace_a",
-        "job_a",
-        DEFAULT_DEPTH,
-        DEFAULT_WITH_DOWNSTREAM
+        namespace="namespace_a", job="job_a", depth=DEFAULT_DEPTH,
+        with_downstream=DEFAULT_WITH_DOWNSTREAM, version="some-version"
     )
 
     assert column_lineage == COLUMN_LINEAGE
-
     mock_get.assert_called_once_with(
-        url=client._url(
-            '/column-lineage'
-        ),
+        url=client._url('/column-lineage'),
         headers=mock.ANY,
         params={
-            'nodeId': 'job:namespace_a:job_a',
+            'nodeId': 'job:namespace_a:job_a#some-version',
             'depth': DEFAULT_DEPTH,
             'withDownstream': DEFAULT_WITH_DOWNSTREAM
         },