Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jaredcatkinson wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions src/openhound_github/graphql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
7 changes: 2 additions & 5 deletions src/openhound_github/helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__)


Expand Down Expand Up @@ -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:
Expand Down
140 changes: 84 additions & 56 deletions src/openhound_github/resources/enterprise.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from openhound_github.graphql import (
ENTERPRISE_ADMINS_QUERY,
ENTERPRISE_MEMBERS_QUERY,
ENTERPRISE_SAML_PROVIDER_QUERY,
ENTERPRISE_QUERY,
ENTERPRISE_SAML_QUERY,
)
Expand Down Expand Up @@ -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}",
Expand All @@ -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,
}
Comment on lines +70 to +96

@coderabbitai coderabbitai Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Add error handling around the new enterprise_organizations GraphQL pagination call.

Unlike enterprise() (Lines 53-62), which wraps its GraphQL POST in try/except with graceful logging, this newly-added independent pagination call has no error handling. Since client.paginate()/.post() raise HTTPError on non-2xx responses by default, a transient or permission error here will propagate as an unhandled exception and fail this resource entirely instead of degrading gracefully.

🛡️ Suggested fix 
 def enterprise_organizations(enterprise_data: 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},
     }

-    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,
-                }
+    try:
+        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,
+                    }
+    except Exception as e:
+        logger.error(
+            f"Error in resource 'enterprise_organizations' processing enterprise '{ctx.enterprise_name}': {e}",
+            extra={"resource": "enterprise_organizations", "phase": "resource_iteration"},
+        )
+        return
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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,
}
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",
):
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,
}
except Exception as e:
logger.error(
f"Error in resource 'enterprise_organizations' processing enterprise '{ctx.enterprise_name}': {e}",
extra={"resource": "enterprise_organizations", "phase": "resource_iteration"},
)
return
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/openhound_github/resources/enterprise.py` around lines 70 - 96, Add the
same graceful error handling used in enterprise() around the new
enterprise_organizations GraphQL pagination flow. Wrap the
ctx.client.paginate("/graphql", ...) loop in try/except (including HTTPError or
a broader request exception if consistent with the module), log the failure with
enough context, and return/stop cleanly instead of letting the exception escape.
Use the enterprise_organizations paginator setup and the ctx.client.paginate
call as the key locations to update.



@app.transformer(name="enterprise_members", columns=BaseUser, parallelized=True)
Expand Down Expand Up @@ -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,
}
Comment on lines +345 to +372

@coderabbitai coderabbitai Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Missing error handling around the SSO-backed GraphQL POST.

client.post("/graphql", json=data).json() at Line 358 is unwrapped, unlike enterprise()'s equivalent call in the same file/diff which catches exceptions and logs gracefully (Lines 53-62). A non-2xx response (e.g., insufficient PAT scope, transient network failure) raises HTTPError by default and will crash this resource instead of degrading gracefully like its sibling.

🛡️ Suggested fix 
-    response = client.post("/graphql", json=data).json()
-    enterprise_object = (response.get("data") or {}).get("enterprise", {})
-    saml_provider = (enterprise_object.get("ownerInfo") or {}).get("samlIdentityProvider")
+    try:
+        response = client.post("/graphql", json=data).json()
+    except Exception as e:
+        logger.error(
+            f"Error in resource 'enterprise_saml_provider' processing enterprise '{ctx.enterprise_name}': {e}",
+            extra={"resource": "enterprise_saml_provider", "phase": "resource_iteration"},
+        )
+        return
+    enterprise_object = (response.get("data") or {}).get("enterprise", {})
+    saml_provider = (enterprise_object.get("ownerInfo") or {}).get("samlIdentityProvider")
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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,
}
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_PROVIDER_QUERY,
"variables": {"slug": ctx.enterprise_name},
}
try:
response = client.post("/graphql", json=data).json()
except Exception as e:
logger.error(
f"Error in resource 'enterprise_saml_provider' processing enterprise '{ctx.enterprise_name}': {e}",
extra={"resource": "enterprise_saml_provider", "phase": "resource_iteration"},
)
return
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,
}
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/openhound_github/resources/enterprise.py` around lines 345 - 372, Add
error handling around the SSO GraphQL request in enterprise_saml_provider: the
unguarded client.post("/graphql", json=data).json() call can raise on non-2xx or
network failures and should degrade gracefully like enterprise(). Wrap the
POST/JSON fetch in a try/except, log a warning or error with ctx.enterprise_name
and the exception details, then return early when the SSO-backed lookup fails;
keep the existing saml_provider parsing and yield logic unchanged when the
request succeeds.



@app.transformer(
Expand All @@ -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",
Expand All @@ -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,
Expand All @@ -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"
Expand All @@ -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),
Expand All @@ -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)
7 changes: 6 additions & 1 deletion src/openhound_github/source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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={
Expand All @@ -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,
Expand Down
Loading