88
99from azure .cli .core .azclierror import AzureResponseError , ResourceNotFoundError
1010from azure .core .exceptions import HttpResponseError
11+ from azure .core .rest import HttpRequest
1112from knack .log import get_logger
1213from rich .console import Console
1314
1415from azext_iot ._factory import adr_service_factory
16+ from azext_iot .constants import LRO_POLL_RETRIES , LRO_POLL_WAIT_SEC
1517from azext_iot .common .utility import wait_for_terminal_state
1618
1719__all__ = ["ADRProvider" , "console" ]
2325console = Console ()
2426
2527
28+ # ---------------------------------------------------------------------------
29+ # TEMPORARY WORKAROUND — the Device Registry service is under active development.
30+ #
31+ # ADR mutations are long-running operations (LROs). The status URL the service
32+ # returns (via the ``Azure-AsyncOperation`` header) points at the DIRECT
33+ # resource-provider host (``control-plane.prod.<region>.iotadr.net``), NOT at
34+ # ARM. That host currently requires an ARM PoP (proof-of-possession) token that
35+ # neither the Azure SDK poller nor a plain bearer client sends, so polling the
36+ # async-operation URL returns HTTP 500 "ARM PoP token authentication failed"
37+ # even though the mutation itself succeeds (HTTP 202). Per the Device Registry
38+ # team this async-status implementation is incomplete/temporary; until it is
39+ # fixed we poll the resource's own ``provisioningState`` instead.
40+ #
41+ # TO REVERT once the backend is fixed: set POLL_PROVISIONING_STATE_WORKAROUND to
42+ # ``False`` (or delete this block, the workaround branch in ``_await_terminal``
43+ # and ``_poll_provisioning_state``). No command call sites need to change.
44+ # ---------------------------------------------------------------------------
45+ POLL_PROVISIONING_STATE_WORKAROUND = True
46+
47+ # Terminal ARM ``provisioningState`` values.
48+ _PROVISIONING_SUCCEEDED = "Succeeded"
49+ _PROVISIONING_FAILURES = ("Failed" , "Canceled" )
50+
51+ # Only PUT/PATCH/DELETE LROs address the resource itself, so their initial
52+ # request URL can be GET-polled for ``provisioningState``. POST-style action
53+ # LROs (synchronize, revoke, refresh, run, ...) have no such pollable resource
54+ # URL, so they fall back to the default SDK polling.
55+ _RESOURCE_MUTATION_METHODS = ("PUT" , "PATCH" , "DELETE" )
56+
57+
2658class ADRProvider (object ):
2759 def __init__ (self , cmd ):
2860 self .cmd = cmd
@@ -35,13 +67,106 @@ def _wait(self, poller, status_message: str, **kwargs):
3567 pop ``no_wait`` from kwargs and return the poller immediately when set,
3668 otherwise show ``status_message`` in a spinner while waiting for the
3769 operation to reach a terminal state. Remaining kwargs are forwarded to
38- ``wait_for_terminal_state `` (e.g. polling interval overrides).
70+ ``_await_terminal `` (e.g. polling interval overrides).
3971 """
4072 no_wait = kwargs .pop ("no_wait" , False )
4173 if no_wait :
4274 return poller
4375 with console .status (status_message ):
44- return wait_for_terminal_state (poller , ** kwargs )
76+ return self ._await_terminal (poller , ** kwargs )
77+
78+ def _await_terminal (self , poller , ** kwargs ):
79+ """Wait for an ADR long-running operation to reach a terminal state.
80+
81+ Single choke point used both by ``_wait`` and by the few commands that
82+ drive their poller directly. While the temporary workaround is enabled we
83+ poll the resource's ``provisioningState`` (see the module note above);
84+ otherwise we defer to the shared SDK helper.
85+ """
86+ if POLL_PROVISIONING_STATE_WORKAROUND :
87+ return self ._poll_provisioning_state (poller , ** kwargs )
88+ return wait_for_terminal_state (poller , ** kwargs )
89+
90+ @staticmethod
91+ def _poller_initial_request (poller ):
92+ """Best-effort ``(url, method)`` of the mutating request a poller wraps.
93+
94+ That request URL *is* the resource URL, which is exactly what we GET to
95+ read ``provisioningState``. Reaches through SDK internals defensively and
96+ returns ``(None, None)`` when the poller is not shaped as expected, so the
97+ caller can fall back to the default polling.
98+ """
99+ method = getattr (poller , "_polling_method" , None )
100+ if method is None and hasattr (poller , "polling_method" ):
101+ # Best-effort only; any failure just falls back to default polling.
102+ try :
103+ method = poller .polling_method ()
104+ except Exception : # noqa: BLE001
105+ method = None
106+ initial = getattr (method , "_initial_response" , None )
107+ request = getattr (initial , "http_request" , None )
108+ if request is None :
109+ request = getattr (getattr (initial , "http_response" , None ), "request" , None )
110+ if request is None :
111+ return None , None
112+ return getattr (request , "url" , None ), (getattr (request , "method" , "" ) or "" ).upper ()
113+
114+ def _poll_provisioning_state (self , poller , wait_sec : int = LRO_POLL_WAIT_SEC , ** _ ):
115+ """TEMPORARY: resolve an LRO by polling the resource's ``provisioningState``.
116+
117+ See the module-level workaround note. Steps:
118+
119+ * If the poller already completed inline (e.g. a create/PUT that returns
120+ the resource with a terminal ``provisioningState``), return its result
121+ directly - no async-operation/PoP hop needed.
122+ * Otherwise recover the resource URL from the poller and GET it until
123+ ``provisioningState`` is terminal: ``Succeeded`` (or a resource that
124+ exposes no ``provisioningState``) returns the body; ``Failed`` /
125+ ``Canceled`` raises ``AzureResponseError``.
126+
127+ For a DELETE a 404 means the resource is gone (success). For any other
128+ method a 404 is treated as "not readable yet" and retried (the by-name
129+ read path can briefly 404 on some backend partitions). Action-style
130+ (POST) LROs and any poller we cannot introspect fall back to the default
131+ SDK polling, so behavior is never worse than before the workaround.
132+ """
133+ from time import sleep
134+
135+ # Fast path: an LRO that completed inline is already terminal with no
136+ # further network calls, so return its result and skip the PoP-500 hop.
137+ if poller .done ():
138+ return poller .result ()
139+
140+ url , method = self ._poller_initial_request (poller )
141+ if not url or method not in _RESOURCE_MUTATION_METHODS :
142+ # No pollable resource URL (e.g. a POST action) -> default polling.
143+ return wait_for_terminal_state (poller , wait_sec = wait_sec )
144+
145+ is_delete = method == "DELETE"
146+ last_body = None
147+ for _ in range (LRO_POLL_RETRIES ):
148+ response = self .client .send_request (HttpRequest ("GET" , url ))
149+ code = response .status_code
150+ if code == 404 :
151+ if is_delete :
152+ return None # resource removed -> delete complete
153+ sleep (wait_sec ) # not readable yet -> retry
154+ continue
155+ if 200 <= code < 300 :
156+ last_body = response .json ()
157+ state = (last_body .get ("properties" ) or {}).get ("provisioningState" )
158+ if state == _PROVISIONING_SUCCEEDED or state is None :
159+ return last_body
160+ if state in _PROVISIONING_FAILURES :
161+ raise AzureResponseError (
162+ f"The operation did not succeed (provisioningState='{ state }
163+ "Inspect the service activity log using the operation's correlation id."
164+ )
165+ sleep (wait_sec ) # still provisioning (Accepted/Updating/...) -> re-check
166+ continue
167+ # Unexpected status (e.g. a transient 5xx) -> brief backoff and retry.
168+ sleep (wait_sec )
169+ return last_body # budget exhausted -> return last-known body (may be None)
45170
46171 def _raise_if_parent_not_found (self , error : Exception , message : str ):
47172 """Translate a backend "ParentResourceNotFound" 404 into a friendly error.
0 commit comments