BED-8788: add PAT-backed enterprise SSO collection and fix repeated enterprise SSO pagination

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) {

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:

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)

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,