Merge to 14 (#15)

* fix: handle missing key more elegantly (#11)

* fix: handle missing key more elegantly

* fix: simplify type

* Update version

* Allow skipping the extraction of organization members and/or repositories. (#13)

* feat: allow skipping org repos or members

* chore: refactor args

* fix: user-repo gathering should only import owned (#9)

* fix: user-repo gathering should only import owned

* Upgrade to nodestream 0.14.0 (#10)

* chore: upgrade to newest nodestream

* chore: bump version

* chore: merge/rebase

* fix: tests

* chore: poetry update

Author: jbristow

nodestream_github/client/githubclient.py

@@ -23,7 +23,7 @@
 
 import nodestream_github.types as types
 from nodestream_github.logging import get_plugin_logger
-from nodestream_github.types.enums import CollaboratorAffiliation
+from nodestream_github.types import enums
 
 DEFAULT_REQUEST_RATE_LIMIT_PER_MINUTE = int(13000 / 60)
 DEFAULT_MAX_RETRIES = 20
@@ -261,7 +261,10 @@ async def _get_item(
         return {}
 
     async def fetch_repos_for_org(
-        self, org_login: str, repo_type: str | None = None
+        self,
+        *,
+        org_login: str,
+        repo_type: enums.OrgRepoType | None = None,
     ) -> AsyncGenerator[types.GithubRepo]:
         """Fetches repositories for the specified organization.
 
@@ -287,7 +290,10 @@ async def fetch_repos_for_org(
             _fetch_problem(f"repos for org {org_login}", e)
 
     async def fetch_members_for_org(
-        self, org_login: str, role: str | None = None
+        self,
+        *,
+        org_login: str,
+        role: enums.OrgMemberRole | None = None,
     ) -> AsyncGenerator[types.GithubUser]:
         """Fetch all users who are members of an organization.
 
@@ -339,8 +345,9 @@ async def fetch_full_org(self, org_login: str) -> types.GithubOrg | None:
 
     async def fetch_repos_for_user(
         self,
+        *,
         user_login: str,
-        repo_type: str | None = None,
+        repo_type: enums.UserRepoType | None = None,
     ) -> AsyncGenerator[types.GithubRepo]:
         """Fetches repositories for a user.
 
@@ -361,7 +368,10 @@ async def fetch_repos_for_user(
             _fetch_problem(f"repos for user {user_login}", e)
 
     async def fetch_languages_for_repo(
-        self, owner_login: str, repo_name: str
+        self,
+        *,
+        owner_login: str,
+        repo_name: str,
     ) -> AsyncGenerator[str]:
         """Fetch languages for the specified repository.
 
@@ -380,7 +390,10 @@ async def fetch_languages_for_repo(
             _fetch_problem(f"languages for repo {owner_login}/{repo_name}", e)
 
     async def fetch_webhooks_for_repo(
-        self, owner_login: str, repo_name: str
+        self,
+        *,
+        owner_login: str,
+        repo_name: str,
     ) -> AsyncGenerator[types.Webhook]:
         """Try to get types.webhook data for this repo.
 
@@ -402,7 +415,7 @@ async def fetch_collaborators_for_repo(
         *,
         owner_login: str,
         repo_name: str,
-        affiliation: CollaboratorAffiliation,
+        affiliation: enums.CollaboratorAffiliation,
     ) -> AsyncGenerator[types.GithubUser]:
         """Try to get collaborator data for this repo.
 
@@ -469,7 +482,9 @@ async def fetch_all_users(self) -> AsyncGenerator[types.GithubUser]:
             _fetch_problem("all users", e)
 
     async def fetch_teams_for_org(
-        self, org_login: str
+        self,
+        *,
+        org_login: str,
     ) -> AsyncGenerator[types.GithubTeamSummary]:
         """Fetch all teams in an organization visible to the authenticated user.
 
@@ -487,7 +502,7 @@ async def fetch_teams_for_org(
         except httpx.HTTPError as e:
             _fetch_problem(f"teams for org {org_login}", e)
 
-    async def fetch_team(self, org_login: str, slug: str) -> types.GithubTeam | None:
+    async def fetch_team(self, *, org_login: str, slug: str) -> types.GithubTeam | None:
         """Fetches a single team for an org by the team slug.
 
         https://docs.github.com/en/[email protected]/rest/teams/teams?apiVersion=2022-11-28#get-a-team-by-name
@@ -498,7 +513,10 @@ async def fetch_team(self, org_login: str, slug: str) -> types.GithubTeam | None
             _fetch_problem(f"full team info for {org_login}/{slug}", e)
 
     async def fetch_members_for_team(
-        self, team_id: int, role: str | None = None
+        self,
+        *,
+        team_id: int,
+        role: enums.TeamMemberRole | None = None,
     ) -> AsyncGenerator[types.GithubUser]:
         """Fetch all users that have a given role for a specified team.
 
@@ -523,7 +541,10 @@ async def fetch_members_for_team(
             _fetch_problem(f"members for team {team_id}", e)
 
     async def fetch_repos_for_team(
-        self, org_login: str, slug: str
+        self,
+        *,
+        org_login: str,
+        slug: str,
     ) -> AsyncGenerator[types.GithubRepo]:
         """Fetch all repos for a specified team visible to the authenticated user.
 

nodestream_github/orgs.py

@@ -14,13 +14,24 @@
 from .interpretations.relationship.user import simplify_user
 from .logging import get_plugin_logger
 from .types import OrgRecord, SimplifiedUser
+from .types.enums import OrgMemberRole
 
 logger = get_plugin_logger(__name__)
 
 
 class GithubOrganizationsExtractor(Extractor):
-    def __init__(self, **github_client_kwargs: any):
-        self.client = GithubRestApiClient(**github_client_kwargs)
+    def __init__(
+        self,
+        *,
+        include_members: bool | None = True,
+        include_repositories: bool | None = True,
+        **kwargs: any,
+    ):
+
+        self.include_members = include_members is True
+        self.include_repositories = include_repositories is True
+
+        self.client = GithubRestApiClient(**kwargs)
 
     async def extract_records(self) -> AsyncGenerator[OrgRecord]:
         async for org in self.client.fetch_all_organizations():
@@ -34,19 +45,34 @@ async def _extract_organization(self, login: str) -> OrgRecord | None:
         if not full_org:
             return None
 
-        full_org["members"] = [user async for user in self._fetch_all_members(login)]
+        if self.include_members:
+            full_org["members"] = [
+                user async for user in self._fetch_all_members(login)
+            ]
+        else:
+            full_org["members"] = []
+
+        if self.include_repositories:
+            full_org["repositories"] = [
+                simplify_repo(
+                    repo, permission=full_org.get("default_repository_permission")
+                )
+                async for repo in self.client.fetch_repos_for_org(org_login=login)
+            ]
+        else:
+            full_org["repositories"] = []
 
-        full_org["repositories"] = [
-            simplify_repo(
-                repo, permission=full_org.get("default_repository_permission")
-            )
-            async for repo in self.client.fetch_repos_for_org(login)
-        ]
         return full_org
 
     async def _fetch_all_members(self, login: str) -> AsyncGenerator[SimplifiedUser]:
-        async for admin in self.client.fetch_members_for_org(login, "admin"):
+        async for admin in self.client.fetch_members_for_org(
+            org_login=login,
+            role=OrgMemberRole.ADMIN,
+        ):
             yield simplify_user(admin) | {"role": "admin"}
 
-        async for member in self.client.fetch_members_for_org(login, "member"):
+        async for member in self.client.fetch_members_for_org(
+            org_login=login,
+            role=OrgMemberRole.MEMBER,
+        ):
             yield simplify_user(member) | {"role": "member"}

nodestream_github/repos.py

@@ -14,7 +14,7 @@
 from .interpretations.relationship.user import simplify_user
 from .logging import get_plugin_logger
 from .types import GithubRepo, RepositoryRecord
-from .types.enums import CollaboratorAffiliation
+from .types.enums import CollaboratorAffiliation, OrgRepoType, UserRepoType
 
 logger = get_plugin_logger(__name__)
 
@@ -95,13 +95,15 @@ async def _extract_repo(self, repo: GithubRepo) -> RepositoryRecord:
         repo["languages"] = [
             {"name": lang}
             async for lang in self.client.fetch_languages_for_repo(
-                owner["login"], repo["name"]
+                owner_login=owner["login"],
+                repo_name=repo["name"],
             )
         ]
         repo["webhooks"] = [
             hook
             async for hook in self.client.fetch_webhooks_for_repo(
-                owner["login"], repo["name"]
+                owner_login=owner["login"],
+                repo_name=repo["name"],
             )
         ]
         repo["collaborators"] = []
@@ -126,12 +128,14 @@ async def _fetch_repos_by_org(self) -> AsyncGenerator[GithubRepo]:
         async for org in self.client.fetch_all_organizations():
             if self.collecting.org_public:
                 async for repo in self.client.fetch_repos_for_org(
-                    org["login"], "public"
+                    org_login=org["login"],
+                    repo_type=OrgRepoType.PUBLIC,
                 ):
                     yield repo
             if self.collecting.org_private:
                 async for repo in self.client.fetch_repos_for_org(
-                    org["login"], "private"
+                    org_login=org["login"],
+                    repo_type=OrgRepoType.PRIVATE,
                 ):
                     yield repo
 
@@ -144,13 +148,8 @@ async def _fetch_repos_by_user(self) -> AsyncGenerator[GithubRepo]:
         repository permissions (read)
         """
         async for user in self.client.fetch_all_users():
-            if self.collecting.user_public:
-                async for repo in self.client.fetch_repos_for_user(
-                    user["login"], "public"
-                ):
-                    yield repo
-            if self.collecting.user_private:
-                async for repo in self.client.fetch_repos_for_user(
-                    user["login"], "private"
-                ):
-                    yield repo
+            async for repo in self.client.fetch_repos_for_user(
+                user_login=user["login"],
+                repo_type=UserRepoType.OWNER,
+            ):
+                yield repo

nodestream_github/teams.py

@@ -14,6 +14,7 @@
 from .interpretations.relationship.user import simplify_user
 from .logging import get_plugin_logger
 from .types import GithubTeam, GithubTeamSummary, SimplifiedUser, TeamRecord
+from .types.enums import TeamMemberRole
 
 logger = get_plugin_logger(__name__)
 
@@ -25,7 +26,7 @@ def __init__(self, **github_client_kwargs: any):
     async def extract_records(self) -> AsyncGenerator[TeamRecord]:
         async for page in self.client.fetch_all_organizations():
             login = page["login"]
-            async for team in self.client.fetch_teams_for_org(login):
+            async for team in self.client.fetch_teams_for_org(org_login=login):
                 team_record = await self._fetch_team(login, team)
                 if team_record:
                     logger.debug(
@@ -42,24 +43,31 @@ async def _fetch_members(self, team: GithubTeam) -> AsyncGenerator[SimplifiedUse
             team["slug"],
         )
 
-        async for member in self.client.fetch_members_for_team(team["id"], "member"):
+        async for member in self.client.fetch_members_for_team(
+            team_id=team["id"],
+            role=TeamMemberRole.MEMBER,
+        ):
             yield member | {"role": "member"}
         async for member in self.client.fetch_members_for_team(
-            team["id"], "maintainer"
+            team_id=team["id"],
+            role=TeamMemberRole.MAINTAINER,
         ):
             yield member | {"role": "maintainer"}
 
     async def _fetch_team(
         self, login: str, team_summary: GithubTeamSummary
     ) -> GithubTeam:
-        team = await self.client.fetch_team(login, team_summary["slug"])
+        team = await self.client.fetch_team(org_login=login, slug=team_summary["slug"])
         if not team:
             return None
         team["members"] = [
             simplify_user(member) async for member in self._fetch_members(team)
         ]
         team["repos"] = [
             simplify_repo(repo, permission=team.get("permission"))
-            async for repo in self.client.fetch_repos_for_team(login, team["slug"])
+            async for repo in self.client.fetch_repos_for_team(
+                org_login=login,
+                slug=team["slug"],
+            )
         ]
         return team

nodestream_github/transformer/repo.py

@@ -19,19 +19,33 @@ def __init__(
         full_name_key: str = "full_name",
         **kwargs: any,
     ):
-
         self.client = GithubRestApiClient(**kwargs)
         self.full_name_key = full_name_key
 
     async def transform_record(
         self,
         record: types.GithubRepo,
     ) -> AsyncGenerator[types.GithubUser]:
-        (repo_owner, repo_name) = record[self.full_name_key].split("/")
-        logging.debug("Transforming record %s/%s", repo_owner, repo_name)
+        logging.debug("Attempting to transform %s", record)
 
+        full_name = record.get(self.full_name_key)
         simplified_repo = simplify_repo(record)
 
+        if full_name is not None:
+            async for user in self._transform(full_name, simplified_repo):
+                yield user
+        else:
+            logging.info("No full_name key found in record %s", record)
+
+    async def _transform(
+        self,
+        full_name: str,
+        simplified_repo: types.SimplifiedRepo,
+    ) -> AsyncGenerator[types.GithubUser]:
+        (repo_owner, repo_name) = full_name.split("/")
+
+        logging.debug("Transforming repo %s/%s", repo_owner, repo_name)
+
         async for collaborator in self.client.fetch_collaborators_for_repo(
             owner_login=repo_owner,
             repo_name=repo_name,

nodestream_github/types/enums.py

@@ -5,3 +5,31 @@ class CollaboratorAffiliation(StrEnum):
     ALL = "all"
     OUTSIDE = "outside"
     DIRECT = "direct"
+
+
+class OrgRepoType(StrEnum):
+    ALL = "all"
+    PUBLIC = "public"
+    PRIVATE = "private"
+    FORKS = "forks"
+    SOURCES = "sources"
+    MEMBER = "member"
+    INTERNAL = "internal"
+
+
+class UserRepoType(StrEnum):
+    ALL = "all"
+    OWNER = "owner"
+    MEMBER = "member"
+
+
+class OrgMemberRole:
+    ALL = "all"
+    ADMIN = "admin"
+    MEMBER = "member"
+
+
+class TeamMemberRole:
+    ALL = "all"
+    MAINTAINER = "maintainer"
+    MEMBER = "member"

nodestream_github/users.py

@@ -13,6 +13,7 @@
 from .interpretations.relationship.repository import simplify_repo
 from .logging import get_plugin_logger
 from .types import UserRecord
+from .types.enums import UserRepoType
 
 logger = get_plugin_logger(__name__)
 
@@ -27,7 +28,10 @@ async def extract_records(self) -> AsyncGenerator[UserRecord]:
             login = user["login"]
             user["repositories"] = [
                 simplify_repo(repo)
-                async for repo in self.client.fetch_repos_for_user(login, "all")
+                async for repo in self.client.fetch_repos_for_user(
+                    user_login=login,
+                    repo_type=UserRepoType.OWNER,
+                )
             ]
             logger.debug("yielded GithubUser{login=%s}", login)
             yield user