changes to support gateway association/disassociation for api version '2025-08-01-preview' (#17)

Author: bavneetsingh16

.github/CODEOWNERS

@@ -88,7 +88,7 @@
 
 /src/ip-group/ @necusjz @kairu-ms @jsntcy
 
-/src/connectedk8s/ @bavneetsingh16 @deeksha345 @anagg929 @atchutbarli
+/src/connectedk8s/ @bavneetsingh16 @deeksha345 @anagg929 @atchutbarli @bgriddaluru
 
 /src/storagesync/ @jsntcy
 

src/connectedk8s/azext_connectedk8s/_client_factory.py

@@ -27,11 +27,11 @@
 
     from azext_connectedk8s.vendored_sdks import ConnectedKubernetesClient
     from azext_connectedk8s.vendored_sdks.operations import ConnectedClusterOperations
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01 import (
-        ConnectedKubernetesClient as ConnectedKubernetesClient20240701,
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01 import (
+        KubernetesClient as ConnectedKubernetesClient20250801,
     )
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01.operations import (
-        ConnectedClusterOperations as ConnectedClusterOperations20240701,
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01.operations import (
+        ConnectedClusterOperations as ConnectedClusterOperations20250801,
     )
 
 AccessToken = namedtuple("AccessToken", ["token", "expires_on"])
@@ -61,34 +61,46 @@ def cf_connected_cluster(cli_ctx: AzCli, _: Any) -> ConnectedClusterOperations:
     return cf_connectedk8s(cli_ctx).connected_cluster
 
 
-def cf_connectedk8s_prev_2024_07_01(
+def cf_connectedk8s_prev_2025_08_01(
     cli_ctx: AzCli, *_: Any
-) -> ConnectedKubernetesClient20240701:
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01 import (
-        ConnectedKubernetesClient,
+) -> ConnectedKubernetesClient20250801:
+    from azure.core.pipeline.policies import HeadersPolicy
+
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01 import (
+        KubernetesClient,
     )
 
-    client: ConnectedKubernetesClient
+    # Create custom headers policy for PUT requests
+    headers_policy = HeadersPolicy({"x-ms-azurearc-cli": "true"})
+
+    client: KubernetesClient
     access_token = os.getenv(consts.Azure_Access_Token_Variable)
     if access_token is not None:
         validate_custom_token()
         credential = AccessTokenCredential(access_token=access_token)
         client = get_mgmt_service_client(
             cli_ctx,
-            ConnectedKubernetesClient,
+            KubernetesClient,
             subscription_id=os.getenv("AZURE_SUBSCRIPTION_ID"),
             credential=credential,
+            base_url="https://management.azure.com",
+            per_call_policies=[headers_policy],
         )
         return client
 
-    client = get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient)
+    client = get_mgmt_service_client(
+        cli_ctx,
+        KubernetesClient,
+        base_url="https://management.azure.com",
+        per_call_policies=[headers_policy],
+    )
     return client
 
 
-def cf_connected_cluster_prev_2024_07_01(
+def cf_connected_cluster_prev_2025_08_01(
     cli_ctx: AzCli, _: Any
-) -> ConnectedClusterOperations20240701:
-    return cf_connectedk8s_prev_2024_07_01(cli_ctx).connected_cluster
+) -> ConnectedClusterOperations20250801:
+    return cf_connectedk8s_prev_2025_08_01(cli_ctx).connected_cluster
 
 
 def cf_connectedmachine(

src/connectedk8s/azext_connectedk8s/_constants.py

@@ -418,7 +418,7 @@
 
 # Connect Precheck Diagnoser constants
 Cluster_Diagnostic_Checks_Job_Registry_Path = (
-    "azurearck8s/helmchart/stable/clusterdiagnosticchecks:0.2.2"
+    "azurearck8s/helmchart/stable/clusterdiagnosticchecks:1.29.3"
 )
 Cluster_Diagnostic_Checks_Helm_Install_Failed_Fault_Type = (
     "Error while installing cluster diagnostic checks helm release"
@@ -525,3 +525,9 @@
 # "Application code shouldn't block the creation of resources for a resource provider that is in the registering state."
 # See https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/resource-providers-and-types#register-resource-provider
 allowed_rp_registration_states = ["Registering", "Registered"]
+
+GATEWAY_LINK_FAULT_TYPE = "gateway-link-error"
+Gateway_Cluster_Resource_Update_Failed_Fault_Type = (
+    "Gateway-Cluster-Resource-Update-Failed"
+)
+GATEWAY_ASSOCIATE_URL = "https://management.azure.com/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.Kubernetes/connectedClusters/{cluster_name}/providers/Microsoft.HybridCompute/settings/Default?api-version={api_version}"

src/connectedk8s/azext_connectedk8s/_troubleshootutils.py

@@ -24,7 +24,7 @@
 if TYPE_CHECKING:
     from kubernetes.client import AppsV1Api, BatchV1Api, CoreV1Api
 
-    from .vendored_sdks.preview_2024_07_01.models import (
+    from .vendored_sdks.preview_2025_08_01.models import (
         ConnectedCluster,
     )
 
@@ -1000,7 +1000,7 @@ def check_agent_version(
             return consts.Diagnostic_Check_Incomplete
 
         # To get user agent version and the latest agent version
-        user_agent_version = connected_cluster.agent_version  # type: ignore[unreachable]
+        user_agent_version = connected_cluster.agent_version
         current_user_version = user_agent_version.split(".")
         latest_agent_version = azure_arc_agent_version.split(".")
         # Comparing if the user version is compatible or not

src/connectedk8s/azext_connectedk8s/_utils.py

@@ -48,7 +48,7 @@
     from kubernetes.client import CoreV1Api, V1NodeList
     from requests import Response
 
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01.models import (
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01.models import (
         ConnectedCluster,
     )
 
@@ -823,7 +823,7 @@ def get_helm_values(
     chart_location_url = f"{config_dp_endpoint}/{chart_location_url_segment}"
     dp_request_identity = connected_cluster.identity
     identity = connected_cluster.id
-    request_dict = connected_cluster.serialize()
+    request_dict = connected_cluster.as_dict()
     request_dict["identity"]["tenantId"] = dp_request_identity.tenant_id
     request_dict["identity"]["principalId"] = dp_request_identity.principal_id
     request_dict["id"] = identity
@@ -903,6 +903,70 @@ def health_check_dp(cmd: CLICommand, config_dp_endpoint: str) -> bool:
     raise CLIInternalError("Error while performing DP health check")
 
 
+def update_gateway_cluster_link(
+    cmd: CLICommand,
+    subscription_id: str,
+    resource_group: str,
+    cluster_name: str,
+    gateway_resource_id: str | None = None,
+) -> bool:
+    """
+    Associates or disassociates a gateway with a cluster.
+
+    If `gateway_resource_id` is provided, performs association.
+    If `gateway_resource_id` is None, performs disassociation.
+    """
+    api_version = "2025-02-19-preview"
+    is_association = gateway_resource_id is not None
+    resource = cmd.cli_ctx.cloud.endpoints.active_directory_resource_id
+    url = consts.GATEWAY_ASSOCIATE_URL.format(
+        subscription_id=subscription_id,
+        resource_group=resource_group,
+        cluster_name=cluster_name,
+        api_version=api_version,
+    )
+
+    headers = ["Content-Type=application/json", "Accept=application/json"]
+
+    token = os.getenv("AZURE_ACCESS_TOKEN")
+    if token:
+        headers.append(f"Authorization=Bearer {token}")
+
+    operation_type = "association" if is_association else "disassociation"
+    body = {
+        "properties": {
+            "gatewayProperties": {
+                "gatewayResourceId": gateway_resource_id  # None in case of disassociation
+            }
+        }
+    }
+    response = send_request_with_retries(
+        cmd.cli_ctx,
+        method="put",
+        url=url,
+        headers=headers,
+        fault_type=consts.GATEWAY_LINK_FAULT_TYPE,
+        summary=f"Error during gateway {operation_type}",
+        request_body=json.dumps(body),
+        resource=resource,
+    )
+
+    if response.status_code == 200:
+        logger.info(
+            f"Gateway {operation_type} succeeded for cluster '{cluster_name}' in resource group '{resource_group}'."
+        )
+        return True
+
+    telemetry.set_exception(
+        exception=f"Gateway {operation_type} failed",
+        fault_type=consts.GATEWAY_LINK_FAULT_TYPE,
+        summary=f"Gateway {operation_type} failed",
+    )
+    raise CLIInternalError(
+        f"Gateway {operation_type} failed for cluster '{cluster_name}'."
+    )
+
+
 def send_request_with_retries(
     cli_ctx: AzCli,
     method: str,
@@ -992,10 +1056,10 @@ def arm_exception_handler(
         status_code = ex.status_code
         if status_code == 404 and return_if_not_found:
             return
-        if status_code // 100 == 4:
+        if status_code is not None and status_code // 100 == 4:
             telemetry.set_user_fault()
         telemetry.set_exception(exception=ex, fault_type=fault_type, summary=summary)
-        if status_code // 100 == 5:
+        if status_code is not None and status_code // 100 == 5:
             raise AzureInternalError(
                 "Http response error occured while making ARM request: "
                 + str(ex)

src/connectedk8s/azext_connectedk8s/clientproxyhelper/_proxylogic.py

@@ -11,7 +11,7 @@
 import azext_connectedk8s._constants as consts
 import azext_connectedk8s.clientproxyhelper._utils as clientproxyutils
 
-from ..vendored_sdks.models import (
+from ..vendored_sdks.preview_2025_08_01.models import (
     ListClusterUserCredentialProperties,
 )
 
@@ -21,10 +21,10 @@
     from knack.commands import CLICommand
     from requests.models import Response
 
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01.models import (
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01.models import (
         CredentialResults,
     )
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01.operations import (
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01.operations import (
         ConnectedClusterOperations,
     )
 
@@ -83,7 +83,7 @@ def get_cluster_user_credentials(
         authentication_method=auth_method, client_proxy=True
     )
 
-    result: CredentialResults = client.list_cluster_user_credential(  # type: ignore[call-overload]
+    result: CredentialResults = client.list_cluster_user_credential(
         resource_group_name,
         cluster_name,
         list_prop,

src/connectedk8s/azext_connectedk8s/clientproxyhelper/_utils.py

@@ -34,7 +34,7 @@
 
     from knack.commands import CLICommand
 
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01.models import (
+    from azext_connectedk8s.vendored_sdks.preview_2025_08_01.models import (
         CredentialResults,
     )
 
@@ -204,17 +204,31 @@ def prepare_clientproxy_data(response: CredentialResults) -> dict[str, Any]:
     data["kubeconfigs"] = []
     kubeconfig = {}
     kubeconfig["name"] = "Kubeconfig"
-    kubeconfig["value"] = b64encode(response.kubeconfigs[0].value).decode("utf-8")  # type: ignore[index]
+
+    # Check if kubeconfigs exists and has items
+    if response.kubeconfigs and len(response.kubeconfigs) > 0:
+        kubeconfig_value = response.kubeconfigs[0].value
+        if kubeconfig_value is not None:
+            kubeconfig["value"] = b64encode(kubeconfig_value).decode("utf-8")
+
     data["kubeconfigs"].append(kubeconfig)
     data["hybridConnectionConfig"] = {}
-    data["hybridConnectionConfig"]["relay"] = response.hybrid_connection_config.relay  # type: ignore[attr-defined]
-    data["hybridConnectionConfig"]["hybridConnectionName"] = (
-        response.hybrid_connection_config.hybrid_connection_name  # type: ignore[attr-defined]
-    )
-    data["hybridConnectionConfig"]["token"] = response.hybrid_connection_config.token  # type: ignore[attr-defined]
-    data["hybridConnectionConfig"]["expirationTime"] = (
-        response.hybrid_connection_config.expiration_time  # type: ignore[attr-defined]
-    )
+
+    # Check if hybrid_connection_config exists
+    if response.hybrid_connection_config is not None:
+        data["hybridConnectionConfig"]["relay"] = (
+            response.hybrid_connection_config.relay
+        )
+        data["hybridConnectionConfig"]["hybridConnectionName"] = (
+            response.hybrid_connection_config.hybrid_connection_name
+        )
+        data["hybridConnectionConfig"]["token"] = (
+            response.hybrid_connection_config.token
+        )
+        data["hybridConnectionConfig"]["expirationTime"] = (
+            response.hybrid_connection_config.expiration_time
+        )
+
     return data
 
 

src/connectedk8s/azext_connectedk8s/commands.py

@@ -11,8 +11,8 @@
 from azure.cli.core.commands import CliCommandType
 
 from azext_connectedk8s._client_factory import (
-    cf_connected_cluster_prev_2024_07_01,
-    cf_connectedk8s_prev_2024_07_01,
+    cf_connected_cluster_prev_2025_08_01,
+    cf_connectedk8s_prev_2025_08_01,
 )
 
 from ._format import connectedk8s_list_table_format, connectedk8s_show_table_format
@@ -27,12 +27,12 @@ def load_command_table(self: Connectedk8sCommandsLoader, _: list[str] | None) ->
             "azext_connectedk8s.vendored_sdks.preview_2024_07_01.operations#"
             "ConnectedClusterOperations.{}"
         ),
-        client_factory=cf_connectedk8s_prev_2024_07_01,
+        client_factory=cf_connectedk8s_prev_2025_08_01,
     )
     with self.command_group(
         "connectedk8s",
         connectedk8s_sdk,
-        client_factory=cf_connected_cluster_prev_2024_07_01,
+        client_factory=cf_connected_cluster_prev_2025_08_01,
     ) as g:
         g.custom_command("connect", "create_connectedk8s", supports_no_wait=True)
         g.custom_command("update", "update_connected_cluster")