[Client API] Get Control Plane Version (#2874)

* Gate feature by checking backend version Signed-off-by: Future-Outlier <[email protected]> * put context in comments or in serialization settings Signed-off-by: Future-Outlier <[email protected]> * get control plane version and add integration tests Signed-off-by: Future-Outlier <[email protected]> * Update Kevin's advice Signed-off-by: Future-Outlier <[email protected]> --------- Signed-off-by: Future-Outlier <[email protected]>
flyteorg · Nov 5, 2024 · 4bef161 · 4bef161
1 parent 6e4e53b
commit 4bef161
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/flytekit/clients/friendly.py b/flytekit/clients/friendly.py
@@ -11,6 +11,7 @@
 from flyteidl.admin import project_pb2 as _project_pb2
 from flyteidl.admin import task_execution_pb2 as _task_execution_pb2
 from flyteidl.admin import task_pb2 as _task_pb2
+from flyteidl.admin import version_pb2 as _version_pb2
 from flyteidl.admin import workflow_attributes_pb2 as _workflow_attributes_pb2
 from flyteidl.admin import workflow_pb2 as _workflow_pb2
 from flyteidl.core import identifier_pb2 as _identifier_pb2
@@ -1087,3 +1088,16 @@ def get_download_artifact_signed_url(
                 expires_in=expires_in_pb,
             )
         )
+
+    def get_control_plane_version(self) -> str:
+        """
+        Retrieve the Control Plane version from Flyteadmin.
+
+        This method calls Flyteadmin's GetVersion API to obtain the current version information of the control plane.
+        The retrieved version can be used to enable or disable specific features based on the Flyteadmin version.
+
+        Returns:
+            str: The version string of the control plane.
+        """
+        version_response = self._stub.GetVersion(_version_pb2.GetVersionRequest(), metadata=self._metadata)
+        return version_response.control_plane_version.Version
diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py
@@ -762,3 +762,8 @@ def test_register_wf_fast(register):
 def test_fetch_active_launchplan_not_found(register):
     remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
     assert remote.fetch_active_launchplan(name="basic.list_float_wf.fake_wf") is None
+
+def test_get_control_plane_version():
+    client = _SynchronousFlyteClient(PlatformConfig.for_endpoint("localhost:30080", True))
+    version = client.get_control_plane_version()
+    assert version == "unknown" or version.startswith("v")