diff --git a/src/openhound_github/graphql.py b/src/openhound_github/graphql.py index 651b661..cbd1d02 100644 --- a/src/openhound_github/graphql.py +++ b/src/openhound_github/graphql.py @@ -98,7 +98,7 @@ query EnterpriseAdmins($slug: String!, $count: Int = 100, $after: String = null) { enterprise(slug: $slug) { ownerInfo { - admins(first: $count, after: $after) { + admins(first: $count, after: $after, role: OWNER) { edges { node { id @@ -115,6 +115,23 @@ } """ +ORGANIZATION_ENTERPRISE_OWNERS_QUERY = """ +query OrganizationEnterpriseOwners($login: String!, $count: Int = 100, $after: String = null) { + organization(login: $login) { + enterpriseOwners(first: $count, after: $after) { + nodes { + id + login + } + pageInfo { + endCursor + hasNextPage + } + } + } +} +""" + ENTERPRISE_SAML_QUERY = """ query EnterpriseSAML($slug: String!, $count: Int = 100, $after: String = null) { enterprise(slug: $slug) { diff --git a/src/openhound_github/models/enterprise_admin.py b/src/openhound_github/models/enterprise_admin.py index 6afa24c..42177b2 100644 --- a/src/openhound_github/models/enterprise_admin.py +++ b/src/openhound_github/models/enterprise_admin.py @@ -18,4 +18,4 @@ ], ) class EnterpriseAdmin(EnterpriseRoleUser): - pass + role_id: str diff --git a/src/openhound_github/resources/enterprise.py b/src/openhound_github/resources/enterprise.py index 3eadfb6..a9c53c7 100644 --- a/src/openhound_github/resources/enterprise.py +++ b/src/openhound_github/resources/enterprise.py @@ -1,5 +1,6 @@ import logging -from dataclasses import dataclass +from dataclasses import dataclass, field +from typing import Any from dlt.sources.helpers.rest_client.client import RESTClient from dlt.sources.helpers.rest_client.paginators import OffsetPaginator @@ -7,6 +8,7 @@ from openhound_github.graphql import ( ENTERPRISE_ADMINS_QUERY, ENTERPRISE_MEMBERS_QUERY, + ORGANIZATION_ENTERPRISE_OWNERS_QUERY, ENTERPRISE_SAML_PROVIDER_QUERY, ENTERPRISE_QUERY, ENTERPRISE_SAML_QUERY, @@ -65,6 +67,7 @@ class SourceContext: emit_legacy_scim_correlations: bool = False github_deployment_id: str = DEFAULT_GITHUB_DEPLOYMENT_ID github_web_origin: str = DEFAULT_GITHUB_WEB_ORIGIN + organizations: list[Any] = field(default_factory=list) def _graphql_client(ctx: SourceContext) -> tuple[RESTClient, str]: @@ -641,6 +644,60 @@ def enterprise_role_users(role: EnterpriseRole, ctx: SourceContext): @app.transformer(name="enterprise_admins", columns=EnterpriseAdmin, parallelized=True) def enterprise_admins(enterprise_data: Enterprise, ctx: SourceContext): + seen_node_ids: set[str] = set() + owner_info_completed = False + client, graphql_path = _sso_graphql_client(ctx) + if not client: + client, graphql_path = _graphql_client(ctx) + + try: + for row in _enterprise_admins_from_owner_info( + enterprise_data, ctx, client, graphql_path + ): + seen_node_ids.add(row["node_id"]) + yield row + owner_info_completed = True + except Exception as e: + logger.warning( + "Unable to collect enterprise owners from ownerInfo.admins for " + "enterprise '%s'; trying organization.enterpriseOwners fallback: %s", + ctx.enterprise_name, + e, + extra={"resource": "enterprise_admins", "phase": "resource_iteration"}, + ) + + if owner_info_completed and seen_node_ids: + return + + for row in _enterprise_admins_from_organizations(enterprise_data, ctx): + if row["node_id"] in seen_node_ids: + continue + seen_node_ids.add(row["node_id"]) + yield row + + +def _enterprise_admin_row( + node: dict[str, Any], enterprise_data: Enterprise, ctx: SourceContext +) -> dict[str, Any] | None: + node_id = node.get("id") + if not node_id: + return None + return { + "node_id": node_id, + "login": node.get("login"), + "assignment": "direct", + "role_id": "owners", + "enterprise_node_id": enterprise_data.id, + "enterprise_slug": ctx.enterprise_name, + } + + +def _enterprise_admins_from_owner_info( + enterprise_data: Enterprise, + ctx: SourceContext, + client: RESTClient, + graphql_path: str, +): paginator = GraphQLCursorPaginator( page_info_path="data.enterprise.ownerInfo.admins.pageInfo", cursor_variable="after", @@ -652,7 +709,6 @@ def enterprise_admins(enterprise_data: Enterprise, ctx: SourceContext): "query": ENTERPRISE_ADMINS_QUERY, "variables": {"slug": ctx.enterprise_name, "count": 100, "after": None}, } - client, graphql_path = _graphql_client(ctx) for page_data in client.paginate( graphql_path, method="POST", @@ -664,16 +720,60 @@ def enterprise_admins(enterprise_data: Enterprise, ctx: SourceContext): es_data = enterprise_object.get("enterprise", {}) owner_info = es_data.get("ownerInfo") or {} for edge in (owner_info.get("admins") or {}).get("edges") or []: - node = edge.get("node") - if node and node.get("id"): - yield { - "node_id": node["id"], - "login": node.get("login"), - "assignment": "direct", - "role_id": "owners", - "enterprise_node_id": enterprise_data.id, - "enterprise_slug": ctx.enterprise_name, - } + row = _enterprise_admin_row(edge.get("node") or {}, enterprise_data, ctx) + if row: + yield row + + +def _enterprise_admins_from_organizations( + enterprise_data: Enterprise, ctx: SourceContext +): + for org in getattr(ctx, "organizations", None) or []: + paginator = GraphQLCursorPaginator( + page_info_path="data.organization.enterpriseOwners.pageInfo", + cursor_variable="after", + cursor_field="endCursor", + has_next_field="hasNextPage", + allow_missing_page_info=True, + ) + data = { + "query": ORGANIZATION_ENTERPRISE_OWNERS_QUERY, + "variables": {"login": org.org_name, "count": 100, "after": None}, + } + client, graphql_path = graphql_client_and_path( + org.client, getattr(org, "graphql_client", None) + ) + found_owner = False + try: + for page_data in client.paginate( + graphql_path, + method="POST", + json=data, + paginator=paginator, + data_selector="data", + ): + for organization_object in page_data: + organization = organization_object.get("organization", {}) + for node in (organization.get("enterpriseOwners") or {}).get( + "nodes" + ) or []: + row = _enterprise_admin_row(node, enterprise_data, ctx) + if row: + found_owner = True + yield row + except Exception as e: + logger.warning( + "Unable to collect enterprise owners through organization '%s' " + "for enterprise '%s': %s", + org.org_name, + ctx.enterprise_name, + e, + extra={"resource": "enterprise_admins", "phase": "resource_iteration"}, + ) + continue + + if found_owner: + return @app.transformer( diff --git a/tests/test_enterprise_resources.py b/tests/test_enterprise_resources.py index 895d809..e195c78 100644 --- a/tests/test_enterprise_resources.py +++ b/tests/test_enterprise_resources.py @@ -7,6 +7,7 @@ from openhound_github.resources.enterprise import ( SourceContext, enterprise, + enterprise_admins, enterprise_external_identity, enterprise_organizations, enterprise_runner_group_memberships, @@ -50,6 +51,13 @@ def paginate(self, path: str, **kwargs): raise ConnectionError("GraphQL endpoint unreachable") +class _PartiallyFailingPaginateClient(_FakeClient): + def paginate(self, path: str, **kwargs): + self.paginate_calls.append((path, kwargs)) + yield from self.pages + raise ConnectionError("GraphQL endpoint unreachable") + + class _FailingPostClient(_FakeClient): def post(self, path: str, json: dict): self.post_calls.append((path, json)) @@ -220,6 +228,218 @@ def test_enterprise_saml_provider_logs_and_returns_when_provider_is_missing( ) +def test_enterprise_admins_use_pat_backed_owner_info_graphql_client() -> None: + assert enterprise_admins._hints["columns"]["role_id"]["data_type"] == "text" + + app_client = _FakeClient(payload={}, pages=[]) + app_graphql_client = _FakeClient(payload={}, pages=[]) + pat_client = _FakeClient(payload={}, pages=[]) + pat_graphql_client = _FakeClient( + payload={}, + pages=[ + [ + { + "enterprise": { + "ownerInfo": { + "admins": { + "edges": [ + {"node": {"id": "U_1", "login": "alice"}}, + {"node": {"id": "U_2", "login": "bob"}}, + ] + } + } + } + } + ] + ], + ) + ctx = SourceContext( + client=app_client, + graphql_client=app_graphql_client, + sso_client=pat_client, + sso_graphql_client=pat_graphql_client, + enterprise_name="acme", + ) + + rows = list(enterprise_admins.__wrapped__(SimpleNamespace(id="E_1"), ctx)) + + assert rows == [ + { + "node_id": "U_1", + "login": "alice", + "assignment": "direct", + "role_id": "owners", + "enterprise_node_id": "E_1", + "enterprise_slug": "acme", + }, + { + "node_id": "U_2", + "login": "bob", + "assignment": "direct", + "role_id": "owners", + "enterprise_node_id": "E_1", + "enterprise_slug": "acme", + }, + ] + assert pat_graphql_client.paginate_calls[0][0] == "" + assert "role: OWNER" in pat_graphql_client.paginate_calls[0][1]["json"]["query"] + assert app_client.paginate_calls == [] + assert app_graphql_client.paginate_calls == [] + + +def test_enterprise_admins_fall_back_to_org_enterprise_owners_without_pat() -> None: + enterprise_client = _FakeClient(payload={}, pages=[]) + enterprise_graphql_client = _FakeClient( + payload={}, + pages=[[{"enterprise": {"ownerInfo": None}}]], + ) + org_client = _FakeClient(payload={}, pages=[]) + org_graphql_client = _FakeClient( + payload={}, + pages=[ + [ + { + "organization": { + "enterpriseOwners": { + "nodes": [{"id": "U_1", "login": "alice"}] + } + } + } + ] + ], + ) + ctx = SourceContext( + client=enterprise_client, + graphql_client=enterprise_graphql_client, + enterprise_name="acme", + organizations=[ + SimpleNamespace( + org_name="acme-org", + client=org_client, + graphql_client=org_graphql_client, + ) + ], + ) + + rows = list(enterprise_admins.__wrapped__(SimpleNamespace(id="E_1"), ctx)) + + assert rows == [ + { + "node_id": "U_1", + "login": "alice", + "assignment": "direct", + "role_id": "owners", + "enterprise_node_id": "E_1", + "enterprise_slug": "acme", + } + ] + assert enterprise_graphql_client.paginate_calls[0][0] == "" + assert org_graphql_client.paginate_calls[0][0] == "" + assert ( + org_graphql_client.paginate_calls[0][1]["json"]["variables"]["login"] + == "acme-org" + ) + + +def test_enterprise_admins_fall_back_after_owner_info_request_failure(caplog) -> None: + enterprise_graphql_client = _FailingPaginateClient(payload={}) + org_graphql_client = _FakeClient( + payload={}, + pages=[ + [ + { + "organization": { + "enterpriseOwners": { + "nodes": [{"id": "U_1", "login": "alice"}] + } + } + } + ] + ], + ) + ctx = SourceContext( + client=_FakeClient(payload={}), + graphql_client=enterprise_graphql_client, + enterprise_name="acme", + organizations=[ + SimpleNamespace( + org_name="acme-org", + client=_FakeClient(payload={}), + graphql_client=org_graphql_client, + ) + ], + ) + + with caplog.at_level(logging.WARNING, logger="openhound_github.resources.enterprise"): + rows = list(enterprise_admins.__wrapped__(SimpleNamespace(id="E_1"), ctx)) + + assert rows[0]["node_id"] == "U_1" + assert any( + "trying organization.enterpriseOwners fallback" in message + for message in caplog.messages + ) + + +def test_enterprise_admins_fall_back_after_partial_owner_info_pagination_failure( + caplog, +) -> None: + enterprise_graphql_client = _PartiallyFailingPaginateClient( + payload={}, + pages=[ + [ + { + "enterprise": { + "ownerInfo": { + "admins": { + "edges": [{"node": {"id": "U_1", "login": "alice"}}] + } + } + } + } + ] + ], + ) + org_graphql_client = _FakeClient( + payload={}, + pages=[ + [ + { + "organization": { + "enterpriseOwners": { + "nodes": [ + {"id": "U_1", "login": "alice"}, + {"id": "U_2", "login": "bob"}, + ] + } + } + } + ] + ], + ) + ctx = SourceContext( + client=_FakeClient(payload={}), + graphql_client=enterprise_graphql_client, + enterprise_name="acme", + organizations=[ + SimpleNamespace( + org_name="acme-org", + client=_FakeClient(payload={}), + graphql_client=org_graphql_client, + ) + ], + ) + + with caplog.at_level(logging.WARNING, logger="openhound_github.resources.enterprise"): + rows = list(enterprise_admins.__wrapped__(SimpleNamespace(id="E_1"), ctx)) + + assert [row["node_id"] for row in rows] == ["U_1", "U_2"] + assert len(org_graphql_client.paginate_calls) == 1 + assert any( + "trying organization.enterpriseOwners fallback" in message + for message in caplog.messages + ) + + def test_enterprise_external_identity_logs_and_returns_on_pagination_failure( caplog, ) -> None: