Skip to content

Commit cdfd8f2

Browse files
committed
UserAssignedManagedIdentity helpers
1 parent 1169979 commit cdfd8f2

File tree

2 files changed

+65
-31
lines changed

2 files changed

+65
-31
lines changed

msal/imds.py

Lines changed: 62 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -34,18 +34,6 @@ def __init__(self, identifier=None, id_type=None):
3434
class UserAssignedManagedIdentity(ManagedIdentity):
3535
"""Feed an instance of this class to :class:`msal.ManagedIdentityClient`
3636
to acquire token for user-assigned managed identity.
37-
38-
By design, an instance of this class is equivalent to a dict in
39-
one of these shapes::
40-
41-
{"ManagedIdentityIdType": "ClientId", "Id": "foo"}
42-
43-
{"ManagedIdentityIdType": "ResourceId", "Id": "foo"}
44-
45-
{"ManagedIdentityIdType": "ObjectId", "Id": "foo"}
46-
47-
so that you may load it from a json configuration file or an env var,
48-
and feed it to :class:`Client`.
4937
"""
5038
CLIENT_ID = "ClientId"
5139
RESOURCE_ID = "ResourceId"
@@ -56,15 +44,7 @@ class UserAssignedManagedIdentity(ManagedIdentity):
5644
OBJECT_ID: "object_id",
5745
}
5846
def __init__(self, identifier, id_type):
59-
"""Construct a UserAssignedManagedIdentity instance.
60-
61-
:param string identifier: The id.
62-
:param string id_type: It shall be one of these three::
63-
64-
UserAssignedManagedIdentity.CLIENT_ID
65-
UserAssignedManagedIdentity.RESOURCE_ID
66-
UserAssignedManagedIdentity.OBJECT_ID
67-
"""
47+
"""Do not use this contructor. Use the following factory methods instead."""
6848
if id_type not in self._types_mapping:
6949
raise ValueError("id_type only accepts one of: {}".format(
7050
list(self._types_mapping)))
@@ -73,6 +53,36 @@ def __init__(self, identifier, id_type):
7353
id_type=id_type,
7454
)
7555

56+
@classmethod
57+
def from_client_id(cls, identifier):
58+
"""Construct a UserAssignedManagedIdentity instance from a client id.
59+
60+
The outcome will be equivalent to::
61+
62+
{"ManagedIdentityIdType": "ClientId", "Id": "foo"}
63+
"""
64+
return UserAssignedManagedIdentity(identifier, cls.CLIENT_ID)
65+
66+
@classmethod
67+
def from_resource_id(cls, identifier):
68+
"""Construct a UserAssignedManagedIdentity instance from a resource id.
69+
70+
The outcome will be equivalent to::
71+
72+
{"ManagedIdentityIdType": "ResourceId", "Id": "foo"}
73+
"""
74+
return UserAssignedManagedIdentity(identifier, cls.RESOURCE_ID)
75+
76+
@classmethod
77+
def from_object_id(cls, identifier):
78+
"""Construct a UserAssignedManagedIdentity instance from an object id.
79+
80+
The outcome will be equivalent to::
81+
82+
{"ManagedIdentityIdType": "ObjectId", "Id": "foo"}
83+
"""
84+
return UserAssignedManagedIdentity(identifier, cls.OBJECT_ID)
85+
7686

7787
class SystemAssignedManagedIdentity(ManagedIdentity):
7888
"""Feed an instance of this class to :class:`msal.ManagedIdentityClient`
@@ -81,9 +91,6 @@ class SystemAssignedManagedIdentity(ManagedIdentity):
8191
By design, an instance of this class is equivalent to::
8292
8393
{"ManagedIdentityIdType": "SystemAssignedManagedIdentity", "Id": None}
84-
85-
so that you may load it from a json configuration file or an env var,
86-
and feed it to :class:`Client`.
8794
"""
8895
def __init__(self):
8996
super(SystemAssignedManagedIdentity, self).__init__(
@@ -257,16 +264,46 @@ def __init__(self, http_client, managed_identity, token_cache=None):
257264
258265
:param dict managed_identity:
259266
It accepts an instance of :class:`SystemAssignedManagedIdentity`
260-
or :class:`UserAssignedManagedIdentity`, or their equivalent dict.
267+
or :class:`UserAssignedManagedIdentity`.
268+
They are equivalent to a dict with a certain shape,
269+
which may be loaded from a json configuration file or an env var,
261270
262271
:param token_cache:
263272
Optional. It accepts a :class:`msal.TokenCache` instance to store tokens.
273+
274+
Example: Hard code a managed identity for your app::
275+
276+
import msal, requests
277+
client = msal.ManagedIdentityClient(
278+
requests.Session(),
279+
msal.UserAssignedManagedIdentity.from_client_id("foo"),
280+
)
281+
282+
Recipe: Write once, run everywhere.
283+
If you use different managed identity on different deployment,
284+
you may use an environment variable (such as AZURE_MANAGED_IDENTITY)
285+
to store a json blob like
286+
``{"ManagedIdentityIdType": "ClientId", "Id": "foo"}`` or
287+
``{"ManagedIdentityIdType": "SystemAssignedManagedIdentity", "Id": null})``.
288+
The following app can load managed identity configuration dynamically::
289+
290+
import json, os, msal, requests
291+
config = os.getenv("AZURE_MANAGED_IDENTITY")
292+
assert config, "An ENV VAR with value should exist"
293+
client = msal.ManagedIdentityClient(
294+
requests.Session(),
295+
json.loads(config),
296+
)
264297
"""
265298
self._http_client = http_client
266299
self._managed_identity = managed_identity
267300
self._token_cache = token_cache
268301

269302
def acquire_token(self, resource=None):
303+
"""Acquire token for the managed identity.
304+
305+
The result will be automatically cached.
306+
"""
270307
if not resource:
271308
raise ValueError(
272309
"The resource parameter is currently required. "

tests/test_mi.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,13 @@
1717
class ManagedIdentityTestCase(unittest.TestCase):
1818
def test_helper_class_should_be_interchangable_with_dict_which_could_be_loaded_from_file_or_env_var(self):
1919
self.assertEqual(
20-
UserAssignedManagedIdentity(
21-
"foo", id_type=UserAssignedManagedIdentity.CLIENT_ID),
20+
UserAssignedManagedIdentity.from_client_id("foo"),
2221
{"ManagedIdentityIdType": "ClientId", "Id": "foo"})
2322
self.assertEqual(
24-
UserAssignedManagedIdentity(
25-
"foo", id_type=UserAssignedManagedIdentity.RESOURCE_ID),
23+
UserAssignedManagedIdentity.from_resource_id("foo"),
2624
{"ManagedIdentityIdType": "ResourceId", "Id": "foo"})
2725
self.assertEqual(
28-
UserAssignedManagedIdentity(
29-
"foo", id_type=UserAssignedManagedIdentity.OBJECT_ID),
26+
UserAssignedManagedIdentity.from_object_id("foo"),
3027
{"ManagedIdentityIdType": "ObjectId", "Id": "foo"})
3128
self.assertEqual(
3229
SystemAssignedManagedIdentity(),

0 commit comments

Comments
 (0)