Repository: gtalarico/pyairtable
Pull request #406 (2024-11-12 06:05:50)

Fix enterprise.info and enterprise.users after new API params
Backport of #405 to 2.x


Patch:
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index ea13684..39c6d46 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+2.3.6 (2024-11-11)
+------------------------
+
+* Fix for `#404 <https://github.com/gtalarico/pyairtable/issues/404>`_
+  related to `enterprise endpoint changes <https://airtable.com/developers/web/api/changelog#anchor-2024-11-11>`__.
+  - `PR #405 <https://github.com/gtalarico/pyairtable/pull/405>`_,
+    `PR #406 <https://github.com/gtalarico/pyairtable/pull/406>`_
+
 2.3.5 (2024-10-29)
 ------------------------
 
diff --git a/pyairtable/__init__.py b/pyairtable/__init__.py
index c895a66..6437f63 100644
--- a/pyairtable/__init__.py
+++ b/pyairtable/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.3.5"
+__version__ = "2.3.6"
 
 from .api import Api, Base, Table
 from .api.enterprise import Enterprise
diff --git a/pyairtable/api/enterprise.py b/pyairtable/api/enterprise.py
index 5eaa711..1cca461 100644
--- a/pyairtable/api/enterprise.py
+++ b/pyairtable/api/enterprise.py
@@ -32,11 +32,26 @@ class Enterprise:
         return self.api.build_url("meta/enterpriseAccounts", self.id)
 
     @cache_unless_forced
-    def info(self) -> EnterpriseInfo:
+    def info(
+        self,
+        *,
+        aggregated: bool = False,
+        descendants: bool = False,
+    ) -> EnterpriseInfo:
         """
         Retrieve basic information about the enterprise, caching the result.
+        Calls `Get enterprise <https://airtable.com/developers/web/api/get-enterprise>`__.
+
+        Args:
+            aggregated: if ``True``, include aggregated values across the enterprise.
+            descendants: if ``True``, include information about the enterprise's descendant orgs.
         """
-        params = {"include": ["collaborators", "inviteLinks"]}
+        include = []
+        if aggregated:
+            include.append("aggregated")
+        if descendants:
+            include.append("descendants")
+        params = {"include": include}
         response = self.api.get(self.url, params=params)
         return EnterpriseInfo.from_api(response, self.api)
 
@@ -54,7 +69,14 @@ class Enterprise:
         payload = self.api.get(url, params=params)
         return UserGroup.parse_obj(payload)
 
-    def user(self, id_or_email: str, collaborations: bool = True) -> UserInfo:
+    def user(
+        self,
+        id_or_email: str,
+        *,
+        collaborations: bool = True,
+        aggregated: bool = False,
+        descendants: bool = False,
+    ) -> UserInfo:
         """
         Retrieve information on a single user with the given ID or email.
 
@@ -62,13 +84,26 @@ class Enterprise:
             id_or_email: A user ID (``usrQBq2RGdihxl3vU``) or email address.
             collaborations: If ``False``, no collaboration data will be requested
                 from Airtable. This may result in faster responses.
+            aggregated: If ``True``, includes the user's aggregated values
+                across this enterprise account and its descendants.
+            descendants: If ``True``, includes information about the user
+                in a ``dict`` keyed per descendant enterprise account.
         """
-        return self.users([id_or_email], collaborations=collaborations)[0]
+        users = self.users(
+            [id_or_email],
+            collaborations=collaborations,
+            aggregated=aggregated,
+            descendants=descendants,
+        )
+        return users[0]
 
     def users(
         self,
         ids_or_emails: Iterable[str],
+        *,
         collaborations: bool = True,
+        aggregated: bool = False,
+        descendants: bool = False,
     ) -> List[UserInfo]:
         """
         Retrieve information on the users with the given IDs or emails.
@@ -80,18 +115,30 @@ class Enterprise:
                 or email addresses (or both).
             collaborations: If ``False``, no collaboration data will be requested
                 from Airtable. This may result in faster responses.
+            aggregated: If ``True``, includes the user's aggregated values
+                across this enterprise account and its descendants.
+            descendants: If ``True``, includes information about the user
+                in a ``dict`` keyed per descendant enterprise account.
         """
         user_ids: List[str] = []
         emails: List[str] = []
         for value in ids_or_emails:
             (emails if "@" in value else user_ids).append(value)
 
+        include = []
+        if collaborations:
+            include.append("collaborations")
+        if aggregated:
+            include.append("aggregated")
+        if descendants:
+            include.append("descendants")
+
         response = self.api.get(
             url=f"{self.url}/users",
             params={
                 "id": user_ids,
                 "email": emails,
-                "include": ["collaborations"] if collaborations else [],
+                "include": include,
             },
         )
         # key by user ID to avoid returning duplicates
diff --git a/pyairtable/models/schema.py b/pyairtable/models/schema.py
index 0de7056..5cb2364 100644
--- a/pyairtable/models/schema.py
+++ b/pyairtable/models/schema.py
@@ -432,11 +432,20 @@ class EnterpriseInfo(AirtableModel):
     user_ids: List[str]
     workspace_ids: List[str]
     email_domains: List["EnterpriseInfo.EmailDomain"]
+    root_enterprise_id: str = pydantic.Field(alias="rootEnterpriseAccountId")
+    descendant_enterprise_ids: List[str] = _FL(alias="descendantEnterpriseAccountIds")
+    aggregated: Optional["EnterpriseInfo.AggregatedIds"] = None
+    descendants: Dict[str, "EnterpriseInfo.AggregatedIds"] = _FD()
 
     class EmailDomain(AirtableModel):
         email_domain: str
         is_sso_required: bool
 
+    class AggregatedIds(AirtableModel):
+        group_ids: List[str] = _FL()
+        user_ids: List[str] = _FL()
+        workspace_ids: List[str] = _FL()
+
 
 class WorkspaceCollaborators(_Collaborators, url="meta/workspaces/{self.id}"):
     """
@@ -566,10 +575,25 @@ class UserInfo(
     is_managed: bool = False
     groups: List[NestedId] = _FL()
     collaborations: "Collaborations" = pydantic.Field(default_factory=Collaborations)
+    descendants: Dict[str, "UserInfo.DescendantIds"] = _FD()
+    aggregated: Optional["UserInfo.AggregatedIds"] = None
 
     def logout(self) -> None:
         self._api.post(self._url + "/logout")
 
+    class DescendantIds(AirtableModel):
+        last_activity_time: Optional[str] = None
+        collaborations: Optional["Collaborations"] = None
+        is_admin: bool = False
+        is_managed: bool = False
+        groups: List[NestedId] = _FL()
+
+    class AggregatedIds(AirtableModel):
+        last_activity_time: Optional[str] = None
+        collaborations: Optional["Collaborations"] = None
+        is_admin: bool = False
+        groups: List[NestedId] = _FL()
+
 
 class UserGroup(AirtableModel):
     """
diff --git a/pyairtable/utils.py b/pyairtable/utils.py
index 90457ea..e03635a 100644
--- a/pyairtable/utils.py
+++ b/pyairtable/utils.py
@@ -210,9 +210,9 @@ def cache_unless_forced(func: Callable[[C], R]) -> FetchMethod[C, R]:
         attr = "_cached_" + attr.lstrip("_")
 
     @wraps(func)
-    def _inner(self: C, *, force: bool = False) -> R:
+    def _inner(self: C, *, force: bool = False, **kwargs: Any) -> R:
         if force or getattr(self, attr, None) is None:
-            setattr(self, attr, func(self))
+            setattr(self, attr, func(self, **kwargs))
         return cast(R, getattr(self, attr))
 
     _inner.__annotations__["force"] = bool


Test patch:
diff --git a/tests/sample_data/EnterpriseInfo.json b/tests/sample_data/EnterpriseInfo.json
index 02ebde5..bd84828 100644
--- a/tests/sample_data/EnterpriseInfo.json
+++ b/tests/sample_data/EnterpriseInfo.json
@@ -1,5 +1,8 @@
 {
   "createdTime": "2019-01-03T12:33:12.421Z",
+  "descendantEnterpriseAccountIds": [
+    "entJ9ZQ5vz9ZQ5vz9"
+  ],
   "emailDomains": [
     {
       "emailDomain": "foobar.com",
@@ -11,6 +14,7 @@
     "ugpR8ZT9KtIgp8Bh3"
   ],
   "id": "entUBq2RGdihxl3vU",
+  "rootEnterpriseAccountId": "entUBq2RGdihxl3vU",
   "userIds": [
     "usrL2PNC5o3H4lBEi",
     "usrsOEchC9xuwRgKk",
diff --git a/tests/sample_data/UserInfo.json b/tests/sample_data/UserInfo.json
index 2097ffe..1edf8f7 100644
--- a/tests/sample_data/UserInfo.json
+++ b/tests/sample_data/UserInfo.json
@@ -12,7 +12,7 @@
       {
         "baseId": "appLkNDICXNqxSDhG",
         "createdTime": "2019-01-03T12:33:12.421Z",
-        "grantedByUserId": "usrqccqnMB2eHylqB",
+        "grantedByUserId": "usrogvSbotRtzdtZW",
         "interfaceId": "pbdyGA3PsOziEHPDE",
         "permissionLevel": "edit"
       }
@@ -38,6 +38,7 @@
   ],
   "id": "usrL2PNC5o3H4lBEi",
   "invitedToAirtableByUserId": "usrsOEchC9xuwRgKk",
+  "isAdmin": true,
   "isManaged": true,
   "isSsoRequired": true,
   "isTwoFactorAuthEnabled": false,
diff --git a/tests/test_api_enterprise.py b/tests/test_api_enterprise.py
index e52c676..0196dc7 100644
--- a/tests/test_api_enterprise.py
+++ b/tests/test_api_enterprise.py
@@ -27,9 +27,10 @@ def enterprise_mocks(enterprise, requests_mock, sample_json):
     m.json_user = sample_json("UserInfo")
     m.json_users = {"users": [m.json_user]}
     m.json_group = sample_json("UserGroup")
+    m.json_enterprise = sample_json("EnterpriseInfo")
     m.user_id = m.json_user["id"]
     m.group_id = m.json_group["id"]
-    m.get_info = requests_mock.get(enterprise.url, json=sample_json("EnterpriseInfo"))
+    m.get_info = requests_mock.get(enterprise.url, json=m.json_enterprise)
     m.get_user = requests_mock.get(
         f"{enterprise.url}/users/{m.user_id}", json=m.json_user
     )
@@ -97,6 +98,35 @@ def test_info(enterprise, enterprise_mocks):
 
     assert enterprise.info(force=True).id == "entUBq2RGdihxl3vU"
     assert enterprise_mocks.get_info.call_count == 2
+    assert "aggregated" not in enterprise_mocks.get_info.last_request.qs
+    assert "descendants" not in enterprise_mocks.get_info.last_request.qs
+
+
+def test_info__aggregated_descendants(enterprise, enterprise_mocks):
+    enterprise_mocks.json_enterprise["aggregated"] = {
+        "groupIds": ["ugp1mKGb3KXUyQfOZ"],
+        "userIds": ["usrL2PNC5o3H4lBEi"],
+        "workspaceIds": ["wspmhESAta6clCCwF"],
+    }
+    enterprise_mocks.json_enterprise["descendants"] = {
+        (sub_ent_id := fake_id("ent")): {
+            "groupIds": ["ugp1mKGb3KXUyDESC"],
+            "userIds": ["usrL2PNC5o3H4DESC"],
+            "workspaceIds": ["wspmhESAta6clDESC"],
+        }
+    }
+    info = enterprise.info(aggregated=True, descendants=True)
+    assert enterprise_mocks.get_info.call_count == 1
+    assert enterprise_mocks.get_info.last_request.qs["include"] == [
+        "aggregated",
+        "descendants",
+    ]
+    assert info.aggregated.group_ids == ["ugp1mKGb3KXUyQfOZ"]
+    assert info.aggregated.user_ids == ["usrL2PNC5o3H4lBEi"]
+    assert info.aggregated.workspace_ids == ["wspmhESAta6clCCwF"]
+    assert info.descendants[sub_ent_id].group_ids == ["ugp1mKGb3KXUyDESC"]
+    assert info.descendants[sub_ent_id].user_ids == ["usrL2PNC5o3H4DESC"]
+    assert info.descendants[sub_ent_id].workspace_ids == ["wspmhESAta6clDESC"]
 
 
 def test_user(enterprise, enterprise_mocks):
@@ -122,6 +152,34 @@ def test_user__no_collaboration(enterprise, enterprise_mocks):
     assert not user.collaborations.workspaces
 
 
+def test_user__descendants(enterprise, enterprise_mocks):
+    enterprise_mocks.json_users["users"][0]["descendants"] = {
+        (other_ent_id := fake_id("ent")): {
+            "lastActivityTime": "2021-01-01T12:34:56Z",
+            "isAdmin": True,
+            "groups": [{"id": (fake_group_id := fake_id("ugp"))}],
+        }
+    }
+    user = enterprise.user(enterprise_mocks.user_id, descendants=True)
+    d = user.descendants[other_ent_id]
+    assert d.last_activity_time == "2021-01-01T12:34:56Z"
+    assert d.is_admin is True
+    assert d.groups[0].id == fake_group_id
+
+
+def test_user__aggregates(enterprise, enterprise_mocks):
+    enterprise_mocks.json_users["users"][0]["aggregated"] = {
+        "lastActivityTime": "2021-01-01T12:34:56Z",
+        "isAdmin": True,
+        "groups": [{"id": (fake_group_id := fake_id("ugp"))}],
+    }
+    user = enterprise.user(enterprise_mocks.user_id, aggregated=True)
+    a = user.aggregated
+    assert a.last_activity_time == "2021-01-01T12:34:56Z"
+    assert a.is_admin is True
+    assert a.groups[0].id == fake_group_id
+
+
 @pytest.mark.parametrize(
     "search_for",
     (
@@ -138,6 +196,36 @@ def test_users(enterprise, search_for):
     assert user.state == "provisioned"
 
 
+def test_users__descendants(enterprise, enterprise_mocks):
+    enterprise_mocks.json_users["users"][0]["descendants"] = {
+        (other_ent_id := fake_id("ent")): {
+            "lastActivityTime": "2021-01-01T12:34:56Z",
+            "isAdmin": True,
+            "groups": [{"id": (fake_group_id := fake_id("ugp"))}],
+        }
+    }
+    users = enterprise.users([enterprise_mocks.user_id], descendants=True)
+    assert len(users) == 1
+    d = users[0].descendants[other_ent_id]
+    assert d.last_activity_time == "2021-01-01T12:34:56Z"
+    assert d.is_admin is True
+    assert d.groups[0].id == fake_group_id
+
+
+def test_users__aggregates(enterprise, enterprise_mocks):
+    enterprise_mocks.json_users["users"][0]["aggregated"] = {
+        "lastActivityTime": "2021-01-01T12:34:56Z",
+        "isAdmin": True,
+        "groups": [{"id": (fake_group_id := fake_id("ugp"))}],
+    }
+    users = enterprise.users([enterprise_mocks.user_id], aggregated=True)
+    assert len(users) == 1
+    a = users[0].aggregated
+    assert a.last_activity_time == "2021-01-01T12:34:56Z"
+    assert a.is_admin is True
+    assert a.groups[0].id == fake_group_id
+
+
 def test_group(enterprise, enterprise_mocks):
     grp = enterprise.group("ugp1mKGb3KXUyQfOZ")
     assert enterprise_mocks.get_group.call_count == 1