From 1fd0bcaf2ae57a8bf8e8802d51f7c1f78c411ae2 Mon Sep 17 00:00:00 2001 From: Jared Atkinson Date: Tue, 7 Jul 2026 08:55:57 -0700 Subject: [PATCH] BED-8788: add PAT-backed enterprise SSO collection and fix repeated enterprise SSO pagination - add optional pat_token support for enterprise app sources - create a dedicated SSO client/context when a PAT is provided - gate enterprise SSO collections on the presence of that client - run enterprise SAML provider and external identity collection with the PAT client - fix enterprise resource iteration so SSO collection does not repeat per org page - give PAT-backed requests the same retry/backoff behavior as the app client - add enterprise resource coverage tests --- src/openhound_github/graphql.py | 20 +++ src/openhound_github/helpers.py | 7 +- src/openhound_github/resources/enterprise.py | 140 +++++++++++-------- src/openhound_github/source.py | 7 +- tests/test_enterprise_resources.py | 93 ++++++++++++ 5 files changed, 205 insertions(+), 62 deletions(-) create mode 100644 tests/test_enterprise_resources.py diff --git a/src/openhound_github/graphql.py b/src/openhound_github/graphql.py index a485f57..e5e4fd1 100644 --- a/src/openhound_github/graphql.py +++ b/src/openhound_github/graphql.py @@ -166,6 +166,26 @@ } """ +ENTERPRISE_SAML_PROVIDER_QUERY = """ +query EnterpriseSAMLProvider($slug: String!) { + enterprise(slug: $slug) { + id + name + slug + ownerInfo { + samlIdentityProvider { + id + issuer + ssoUrl + digestMethod + signatureMethod + idpCertificate + } + } + } +} +""" + TEAMS_QUERY = """ query Teams($login: String!, $count: Int!, $after: String) { organization(login: $login) { diff --git a/src/openhound_github/helpers.py b/src/openhound_github/helpers.py index 78a7c4b..36bd036 100644 --- a/src/openhound_github/helpers.py +++ b/src/openhound_github/helpers.py @@ -4,15 +4,12 @@ from dlt.common import jsonpath from dlt.sources.helpers import requests +from dlt.sources.helpers.rest_client.auth import AuthConfigBase from dlt.sources.helpers.rest_client.paginators import ( JSONResponseCursorPaginator, ) from requests import Request -from openhound_github.auth import ( - GitHubAppInstallationAuth, -) - logger = logging.getLogger(__name__) @@ -154,7 +151,7 @@ def _has_graphql_errors(response: requests.Response) -> bool: return isinstance(response_data, dict) and bool(response_data.get("errors")) -def github_retry_policy(auth: GitHubAppInstallationAuth): +def github_retry_policy(auth: AuthConfigBase): def retry_policy( response: Optional[requests.Response], exception: Optional[BaseException] ) -> bool: diff --git a/src/openhound_github/resources/enterprise.py b/src/openhound_github/resources/enterprise.py index a5d114e..664df2e 100644 --- a/src/openhound_github/resources/enterprise.py +++ b/src/openhound_github/resources/enterprise.py @@ -6,6 +6,7 @@ from openhound_github.graphql import ( ENTERPRISE_ADMINS_QUERY, ENTERPRISE_MEMBERS_QUERY, + ENTERPRISE_SAML_PROVIDER_QUERY, ENTERPRISE_QUERY, ENTERPRISE_SAML_QUERY, ) @@ -37,35 +38,23 @@ class SourceContext: """Shared context for GitHub API access.""" client: RESTClient + sso_client: RESTClient | None = None org_name: str | None = None enterprise_name: str | None = None @app.resource(name="enterprise", columns=Enterprise, parallelized=True) def enterprise(ctx: SourceContext): - paginator = GraphQLCursorPaginator( - page_info_path="data.enterprise.organizations.pageInfo", - cursor_variable="after", - cursor_field="endCursor", - has_next_field="hasNextPage", - ) data = { "query": ENTERPRISE_QUERY, "variables": {"slug": ctx.enterprise_name, "after": None}, } try: - for page_data in ctx.client.paginate( - "/graphql", - method="POST", - json=data, - paginator=paginator, - data_selector="data", - ): - page_enterprise = page_data[0].get("enterprise") - - if page_enterprise: - yield page_enterprise + response = ctx.client.post("/graphql", json=data).json() + page_enterprise = (response.get("data") or {}).get("enterprise") + if page_enterprise: + yield page_enterprise except Exception as e: logger.error( f"Error in resource 'enterprise' processing enterprise '{ctx.enterprise_name}': {e}", @@ -78,13 +67,33 @@ def enterprise(ctx: SourceContext): name="enterprise_organizations", columns=EnterpriseOrganization, parallelized=True ) def enterprise_organizations(enterprise_data: Enterprise, ctx: SourceContext): - orgs = (enterprise_data.organizations or {}).get("nodes", []) - for org in orgs: - yield { - **org, - "enterprise_node_id": enterprise_data.id, - "enterprise_slug": ctx.enterprise_name, - } + paginator = GraphQLCursorPaginator( + page_info_path="data.enterprise.organizations.pageInfo", + cursor_variable="after", + cursor_field="endCursor", + has_next_field="hasNextPage", + ) + data = { + "query": ENTERPRISE_QUERY, + "variables": {"slug": ctx.enterprise_name, "after": None}, + } + + for page_data in ctx.client.paginate( + "/graphql", + method="POST", + json=data, + paginator=paginator, + data_selector="data", + ): + for enterprise_object in page_data: + es_data = enterprise_object.get("enterprise", {}) + orgs = (es_data.get("organizations") or {}).get("nodes", []) + for org in orgs: + yield { + **org, + "enterprise_node_id": enterprise_data.id, + "enterprise_slug": ctx.enterprise_name, + } @app.transformer(name="enterprise_members", columns=BaseUser, parallelized=True) @@ -333,35 +342,34 @@ def enterprise_admins(enterprise_data: Enterprise, ctx: SourceContext): name="enterprise_saml_provider", columns=EnterpriseSamlProvider, parallelized=True ) def enterprise_saml_provider(enterprise_data: Enterprise, ctx: SourceContext): - paginator = GraphQLCursorPaginator( - page_info_path="data.enterprise.ownerInfo.samlIdentityProvider.externalIdentities.pageInfo", - cursor_variable="after", - cursor_field="endCursor", - has_next_field="hasNextPage", - allow_missing_page_info=True, - ) + client = ctx.sso_client + if not client: + logger.info( + "Skipping enterprise_saml_provider for enterprise '%s': no SSO client configured", + ctx.enterprise_name, + ) + return + data = { - "query": ENTERPRISE_SAML_QUERY, - "variables": {"slug": ctx.enterprise_name, "count": 1, "after": None}, + "query": ENTERPRISE_SAML_PROVIDER_QUERY, + "variables": {"slug": ctx.enterprise_name}, } - for page_data in ctx.client.paginate( - "/graphql", - method="POST", - json=data, - paginator=paginator, - data_selector="data", - ): - for enterprise_object in page_data: - es_data = enterprise_object.get("enterprise", {}) - saml_provider = (es_data.get("ownerInfo") or {}).get("samlIdentityProvider") - if not saml_provider: - return - yield { - **{k: v for k, v in saml_provider.items() if k != "externalIdentities"}, - "enterprise_node_id": enterprise_data.id, - "enterprise_slug": ctx.enterprise_name, - } + response = client.post("/graphql", json=data).json() + enterprise_object = (response.get("data") or {}).get("enterprise", {}) + saml_provider = (enterprise_object.get("ownerInfo") or {}).get("samlIdentityProvider") + if not saml_provider: + logger.warning( + "No enterprise SAML provider returned for enterprise '%s'", + ctx.enterprise_name, + ) + return + + yield { + **saml_provider, + "enterprise_node_id": enterprise_data.id, + "enterprise_slug": ctx.enterprise_name, + } @app.transformer( @@ -372,6 +380,14 @@ def enterprise_saml_provider(enterprise_data: Enterprise, ctx: SourceContext): def enterprise_external_identities( saml_provider: EnterpriseSamlProvider, ctx: SourceContext ): + client = ctx.sso_client + if not client: + logger.info( + "Skipping enterprise_external_identities for enterprise '%s': no SSO client configured", + ctx.enterprise_name, + ) + return + paginator = GraphQLCursorPaginator( page_info_path="data.enterprise.ownerInfo.samlIdentityProvider.externalIdentities.pageInfo", cursor_variable="after", @@ -384,7 +400,7 @@ def enterprise_external_identities( "variables": {"slug": ctx.enterprise_name, "count": 100, "after": None}, } - for page_data in ctx.client.paginate( + for page_data in client.paginate( "/graphql", method="POST", json=data, @@ -395,6 +411,10 @@ def enterprise_external_identities( es_data = enterprise_object.get("enterprise", {}) page_provider = (es_data.get("ownerInfo") or {}).get("samlIdentityProvider") if not page_provider: + logger.warning( + "No enterprise SAML provider returned while fetching external identities for enterprise '%s'", + ctx.enterprise_name, + ) return for identity in (page_provider.get("externalIdentities") or {}).get( "nodes" @@ -415,8 +435,7 @@ def enterprise_resources(ctx: SourceContext): members_resource = enterprise_members(ctx) teams_resource = enterprise_teams(ctx) roles_resource = enterprise_roles(ctx) - saml_resource = enterprise_saml_provider(ctx) - return ( + resources = [ enterprise_resource, enterprise_resource | organizations_resource, enterprise_resource | members_resource | enterprise_users(ctx), @@ -430,6 +449,15 @@ def enterprise_resources(ctx: SourceContext): enterprise_resource | roles_resource | enterprise_role_teams(ctx), # enterprise_resource | enterprise_admin_roles(ctx), enterprise_resource | enterprise_admins(ctx), - enterprise_resource | saml_resource, - enterprise_resource | saml_resource | enterprise_external_identities(ctx), - ) + ] + + if ctx.sso_client: + saml_resource = enterprise_saml_provider(ctx) + resources.extend( + [ + enterprise_resource | saml_resource, + enterprise_resource | saml_resource | enterprise_external_identities(ctx), + ] + ) + + return tuple(resources) diff --git a/src/openhound_github/source.py b/src/openhound_github/source.py index 5da35cb..873d1c8 100644 --- a/src/openhound_github/source.py +++ b/src/openhound_github/source.py @@ -8,6 +8,7 @@ from dlt.common.configuration import configspec from dlt.common.configuration.specs import CredentialsConfiguration from dlt.sources.helpers import requests +from dlt.sources.helpers.rest_client.auth import AuthConfigBase from dlt.sources.helpers.rest_client.auth import BearerTokenAuth from dlt.sources.helpers.rest_client.client import RESTClient from dlt.sources.helpers.rest_client.paginators import ( @@ -39,6 +40,7 @@ class OrgContext: class SourceContext: organizations: list[OrgContext] | None = field(default_factory=list) client: RESTClient | None = None + sso_client: RESTClient | None = None enterprise_name: str | None = None cache_lock: Lock = field(default_factory=Lock) app_cache: dict[str, dict[str, Any]] = field(default_factory=dict) @@ -66,6 +68,7 @@ class GithubEnterpriseAppCredentials(CredentialsConfiguration): app_id: str = None key_path: str = None enterprise_name: str = None + pat_token: str | None = None api_uri: str = "https://api.github.com" @property @@ -113,7 +116,7 @@ def source( host (str): The base GitHub API URL used for API calls. """ - def client(auth: GitHubAppInstallationAuth) -> RESTClient: + def client(auth: AuthConfigBase) -> RESTClient: return RESTClient( base_url=host, headers={ @@ -130,6 +133,8 @@ def client(auth: GitHubAppInstallationAuth) -> RESTClient: if credentials.auth == "enterprise_app": ctx = SourceContext(enterprise_name=credentials.enterprise_name) + if credentials.pat_token: + ctx.sso_client = client(BearerTokenAuth(token=credentials.pat_token)) github_app_session = GithubApp( client_id=credentials.client_id, private_key_path=credentials.key_path, diff --git a/tests/test_enterprise_resources.py b/tests/test_enterprise_resources.py new file mode 100644 index 0000000..25cb818 --- /dev/null +++ b/tests/test_enterprise_resources.py @@ -0,0 +1,93 @@ +from types import SimpleNamespace + +from openhound_github.resources.enterprise import ( + SourceContext, + enterprise, + enterprise_organizations, +) + + +class _FakeResponse: + def __init__(self, payload: dict): + self._payload = payload + + def json(self) -> dict: + return self._payload + + +class _FakeClient: + def __init__(self, payload: dict, pages: list[dict] | None = None): + self.payload = payload + self.pages = pages or [] + self.post_calls: list[tuple[str, dict]] = [] + self.paginate_calls: list[tuple[str, dict]] = [] + + def post(self, path: str, json: dict): + self.post_calls.append((path, json)) + return _FakeResponse(self.payload) + + def paginate(self, path: str, **kwargs): + self.paginate_calls.append((path, kwargs)) + return iter(self.pages) + + +def test_enterprise_resource_yields_single_record() -> None: + client = _FakeClient( + { + "data": { + "enterprise": { + "id": "E_1", + "slug": "acme", + "organizations": { + "nodes": [{"id": "O_1", "login": "org-1"}], + "pageInfo": {"hasNextPage": True, "endCursor": "cursor-1"}, + }, + } + } + } + ) + ctx = SourceContext(client=client, enterprise_name="acme") + + rows = list(enterprise(ctx)) + + assert len(rows) == 1 + assert rows[0].id == "E_1" + assert len(client.post_calls) == 1 + + +def test_enterprise_organizations_paginates_all_pages() -> None: + client = _FakeClient( + payload={}, + pages=[ + [ + { + "enterprise": { + "organizations": { + "nodes": [{"id": "O_1", "login": "org-1"}], + "pageInfo": {"hasNextPage": True, "endCursor": "cursor-1"}, + } + } + } + ], + [ + { + "enterprise": { + "organizations": { + "nodes": [{"id": "O_2", "login": "org-2"}], + "pageInfo": {"hasNextPage": False, "endCursor": None}, + } + } + } + ], + ], + ) + ctx = SourceContext(client=client, enterprise_name="acme") + enterprise_data = SimpleNamespace(id="E_1") + + rows = list(enterprise_organizations.__wrapped__(enterprise_data, ctx)) + + assert rows == [ + {"id": "O_1", "login": "org-1", "enterprise_node_id": "E_1", "enterprise_slug": "acme"}, + {"id": "O_2", "login": "org-2", "enterprise_node_id": "E_1", "enterprise_slug": "acme"}, + ] + assert len(client.paginate_calls) == 1