diff --git a/flytekit/clients/friendly.py b/flytekit/clients/friendly.py index 72875f84f7..f506960f82 100644 --- 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 index 4d77e1b610..ee4573c011 100644 --- 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")