Add ability to authenticate multiple orgs in one stream

meltano.yml

@@ -22,7 +22,7 @@ plugins:
     - name: additional_auth_tokens
       kind: array
     - name: auth_app_keys
-      kind: array
+      kind: object
     - name: rate_limit_buffer
       kind: integer
     - name: expiry_time_buffer

tap_github/authenticator.py

@@ -4,6 +4,7 @@
 
 import logging
 import time
+from collections import defaultdict
 from copy import deepcopy
 from datetime import datetime, timedelta, timezone
 from os import environ
@@ -171,6 +172,7 @@ class AppTokenManager(TokenManager):
     def __init__(
         self,
         env_key: str,
+        organization: str | None = None,
         rate_limit_buffer: int | None = None,
         expiry_time_buffer: int | None = None,
         **kwargs,  # noqa: ANN003
@@ -182,6 +184,7 @@ def __init__(
         self.github_app_id = parts[0]
         self.github_private_key = (parts[1:2] or [""])[0].replace("\\n", "\n")
         self.github_installation_id: str | None = parts[2] if len(parts) >= 3 else None
+        self.github_organization: str | None = organization
 
         if expiry_time_buffer is None:
             expiry_time_buffer = self.DEFAULT_EXPIRY_BUFFER_MINS
@@ -246,7 +249,7 @@ class GitHubTokenAuthenticator(APIAuthenticatorBase):
     def get_env():  # noqa: ANN205
         return dict(environ)
 
-    def prepare_tokens(self) -> list[TokenManager]:
+    def prepare_tokens(self) -> dict[str | None, list[TokenManager]]:
         """Prep GitHub tokens"""
 
         env_dict = self.get_env()
@@ -282,39 +285,66 @@ def prepare_tokens(self) -> list[TokenManager]:
                 logger.warning("A token was dismissed.")
 
         # Parse App level private keys and generate tokens
-        # To simplify settings, we use a single env-key formatted as follows:
-        # "{app_id};;{-----BEGIN RSA PRIVATE KEY-----\n_YOUR_PRIVATE_KEY_\n-----END RSA PRIVATE KEY-----}"  # noqa: E501
+        # App keys can be provided via:
+        # 1. Config as array (org-agnostic): ["app_id;;private_key", ...]
+        # 2. Config as dict (org-specific): {"org": ["app_id;;private_key", ...], ...}
+        # 3. Environment variable GITHUB_APP_PRIVATE_KEY: "app_id;;private_key"
+        # Format: "{app_id};;{-----BEGIN RSA PRIVATE KEY-----\n_YOUR_PRIVATE_KEY_\n-----END RSA PRIVATE KEY-----}"  # noqa: E501
 
-        app_keys: set[str] = set()
+        app_keys: dict[str, list[str]] = defaultdict(list)
         if self.auth_app_keys:
-            app_keys = app_keys.union(self.auth_app_keys)
-            logger.info(
-                "Provided %d app keys via config for authentication.",
-                len(app_keys),
-            )
+            # Handle backwards compatibility: detect if it's a list (old format) or dict (new format)
+            if isinstance(self.auth_app_keys, list):
+                # Old format: treat all keys as org-agnostic
+                for app_key in self.auth_app_keys:
+                    app_keys[None].append(app_key)
+                logger.info(
+                    "Provided %d app keys via config for authentication.",
+                    len(self.auth_app_keys),
+                )
+            else:
+                # New format: org-specific keys
+                for org in self.auth_app_keys:
+                    for app_key in self.auth_app_keys[org]:
+                        app_keys[org].append(app_key)
+                    logger.info(
+                      "Provided %d app keys via config for authentication for organization: %s",
+                      len(self.auth_app_keys[org]),
+                      org,
+                    )
         elif "GITHUB_APP_PRIVATE_KEY" in env_dict:
-            app_keys.add(env_dict["GITHUB_APP_PRIVATE_KEY"])
+            app_keys[None].append(env_dict["GITHUB_APP_PRIVATE_KEY"])
             logger.info("Found 1 app key via environment variable for authentication.")
 
-        app_token_managers: list[TokenManager] = []
-        for app_key in app_keys:
-            try:
-                app_token_manager = AppTokenManager(
-                    app_key,
-                    rate_limit_buffer=self.rate_limit_buffer,
-                    expiry_time_buffer=self.expiry_time_buffer,
-                )
-                if app_token_manager.is_valid_token():
-                    app_token_managers.append(app_token_manager)
-            except ValueError as e:  # noqa: PERF203
-                logger.warning(f"An error was thrown while preparing an app token: {e}")
+        token_managers: dict[str | None, list[TokenManager]] = defaultdict(list)
+        for org in app_keys:
+            for app_key in app_keys[org]:
+                try:
+                    app_token_manager = AppTokenManager(
+                        app_key,
+                        organization=org,
+                        rate_limit_buffer=self.rate_limit_buffer,
+                        expiry_time_buffer=self.expiry_time_buffer,
+                    )
+                    if app_token_manager.is_valid_token():
+                        token_managers[org].append(app_token_manager)
+                except ValueError as e:  # noqa: PERF203
+                    logger.warning(
+                        f"An error was thrown while preparing an app token: {e}"
+                    )
 
         logger.info(
             "Tap will run with %d personal auth tokens and %d app keys.",
             len(personal_token_managers),
-            len(app_token_managers),
+            sum([len(token_managers[org]) for org in token_managers]),
         )
-        return personal_token_managers + app_token_managers
+
+        # Merge personal tokens and app tokens.
+        # Personal tokens are stored under None key as they are org-agnostic.
+        if personal_token_managers:
+            token_managers[None].extend(personal_token_managers)
+
+        return token_managers
 
     def __init__(
         self,
@@ -344,9 +374,19 @@ def __init__(
         self.auth_app_keys = auth_app_keys
 
         self.token_managers = self.prepare_tokens()
-        self.active_token: TokenManager | None = (
-            choice(self.token_managers) if self.token_managers else None
-        )
+        self.current_organization: str | None = None
+        if self.token_managers:
+            # Prefer org-specific tokens over org-agnostic (None key)
+            org_keys = [k for k in self.token_managers.keys() if k is not None]
+            if org_keys:
+                initial_org = min(org_keys)
+            else:
+                initial_org = None
+            self.logger.info(f"Setting initial organization for authenticator: {initial_org}")
+            self.active_token: TokenManager | None = choice(self.token_managers[initial_org])
+        else:
+            self.logger.info("Setting initial organization for authenticator: None")
+            self.active_token: TokenManager | None = None
 
     @classmethod
     def from_stream(cls, stream: RESTStream) -> GitHubTokenAuthenticator:
@@ -358,11 +398,69 @@ def from_stream(cls, stream: RESTStream) -> GitHubTokenAuthenticator:
             auth_app_keys=stream.config.get("auth_app_keys"),
         )
 
+    def set_organization(self, org: str) -> None:
+        """Set the current organization and switch to an appropriate token.
+
+        Args:
+            org: The organization name to switch to.
+        """
+        # If we're already using this org, no need to switch
+        if self.current_organization == org:
+            return
+
+        logger.info(f"Switching authentication context to organization: {org}")
+        self.current_organization = org
+
+        # Get tokens for this org (check both org-specific and None keys)
+        available_tokens = self.token_managers.get(org, [])
+        if not available_tokens and None in self.token_managers:
+            # Fall back to org-agnostic tokens (personal tokens or env var app keys)
+            available_tokens = self.token_managers[None]
+            logger.info(
+                f"No org-specific tokens found for '{org}', using org-agnostic tokens"
+            )
+
+        if not available_tokens:
+            logger.warning(
+                f"No authentication tokens available for organization: {org}"
+            )
+            self.active_token = None
+            return
+
+        # Select a token with remaining calls
+        for token_manager in available_tokens:
+            if token_manager.has_calls_remaining():
+                self.active_token = token_manager
+                logger.info(f"Selected token for organization: {org}")
+                return
+
+        # If no tokens have calls remaining, just pick the first one
+        # (it might refresh or we'll rotate later)
+        self.active_token = available_tokens[0]
+        logger.info(
+            f"Selected token for organization: {org} (may need rate limit refresh)"
+        )
+
     def get_next_auth_token(self) -> None:
         current_token = self.active_token.token if self.active_token else ""
-        token_managers = deepcopy(self.token_managers)
-        shuffle(token_managers)
-        for token_manager in token_managers:
+
+        # Build a list of candidate tokens for the current organization
+        candidates = []
+
+        # Priority 1: Other tokens for the current organization
+        if self.current_organization and self.current_organization in self.token_managers:
+            org_tokens = list(self.token_managers[self.current_organization])
+            shuffle(org_tokens)
+            candidates.extend(org_tokens)
+
+        # Priority 2: Org-agnostic tokens (stored under None key)
+        if None in self.token_managers:
+            agnostic_tokens = list(self.token_managers[None])
+            shuffle(agnostic_tokens)
+            candidates.extend(agnostic_tokens)
+
+        # Try to find a token with remaining capacity
+        for token_manager in candidates:
             if (
                 token_manager.has_calls_remaining()
                 and current_token != token_manager.token
@@ -380,7 +478,9 @@ def update_rate_limit(
         response_headers: requests.models.CaseInsensitiveDict,
     ) -> None:
         # If no token or only one token is available, return early.
-        if len(self.token_managers) <= 1 or self.active_token is None:
+        # Count total tokens across all organizations
+        total_tokens = sum(len(tokens) for tokens in self.token_managers.values())
+        if total_tokens <= 1 or self.active_token is None:
             return
 
         self.active_token.update_rate_limit(response_headers)

tap_github/client.py

@@ -237,6 +237,7 @@ def validate_response(self, response: requests.Response) -> None:
                 and "rate limit exceeded" in str(response.content).lower()
             ):
                 # Update token
+                self.logger.info("CONTEXT: " + repr(self.context))
                 self.authenticator.get_next_auth_token()
                 # Raise an error to force a retry with the new token.
                 raise RetriableAPIError(msg, response)

tap_github/organization_streams.py

@@ -70,6 +70,11 @@ def get_records(self, context: Context | None) -> Iterable[dict[str, Any]]:
         quota when only syncing a child stream. Without this,
         the API call is sent but data is discarded.
         """
+        # Set organization-specific authentication before fetching records
+        if context is not None and "org" in context:
+            self.logger.info(f"Setting organization within organization stream: {context['org']}")
+            self.authenticator.set_organization(context["org"])
+
         if (
             not self.selected
             and "skip_parent_streams" in self.config
@@ -653,6 +658,11 @@ def request_records(self, context: Context | None) -> Iterable[dict]:
                     f"Context: {context}. Error: {e}"
                 )
                 return
+            elif "Timeout on validation of query" in error_message:
+                self.logger.warning(
+                    f"Skipping project due to query validation timeout error. "
+                    f"Context: {context}. Error: {e}"
+                )
 
             raise
 

tap_github/repository_streams.py

@@ -89,16 +89,21 @@ class TempStream(GitHubGraphqlStream):
                 th.Property("databaseId", th.IntegerType),
             ).to_dict()
 
-            def __init__(self, tap, repo_list) -> None:  # noqa: ANN001
+            def __init__(self, tap, repo_list, parent_authenticator=None) -> None:  # noqa: ANN001
                 super().__init__(tap)
                 self.repo_list = repo_list
+                # Use parent's authenticator to maintain consistent auth state and rate limits
+                if parent_authenticator is not None:
+                    self._authenticator = parent_authenticator
 
             @property
             def query(self) -> str:
                 chunks = []
                 for i, repo in enumerate(self.repo_list):
+                    org, repo_name = repo
+                    self.logger.info(f"[TempStream.query] Building query for {org}/{repo_name} with authenticator org: {self.authenticator.current_organization}")
                     chunks.append(
-                        f'repo{i}: repository(name: "{repo[1]}", owner: "{repo[0]}") '
+                        f'repo{i}: repository(name: "{repo_name}", owner: "{org}") '
                         "{ nameWithOwner databaseId }"
                     )
                 return "query {" + " ".join(chunks) + " rateLimit { cost } }"
@@ -118,9 +123,20 @@ def validate_response(self, response: requests.Response) -> None:
 
         if len(repo_list) < 1:
             return []
+        self.logger.info(f"Getting repo ids for {len(repo_list)} repositories")
+
+        # Log parent stream authenticator info
+        self.logger.info(f"[get_repo_ids] Parent stream authenticator ID: {id(self.authenticator)}")
+        self.logger.info(f"[get_repo_ids] Parent stream current_org: {self.authenticator.current_organization}")
+        self.logger.info(f"[get_repo_ids] Parent stream active_token org: {getattr(self.authenticator.active_token, 'github_organization', 'N/A')}")
 
         repos_with_ids: list = []
-        temp_stream = TempStream(self._tap, list(repo_list))
+        temp_stream = TempStream(self._tap, list(repo_list), self.authenticator)
+
+        # Log TempStream authenticator info
+        self.logger.info(f"[get_repo_ids] TempStream authenticator ID: {id(temp_stream.authenticator)}")
+        self.logger.info(f"[get_repo_ids] TempStream current_org: {temp_stream.authenticator.current_organization}")
+        self.logger.info(f"[get_repo_ids] TempStream active_token org: {getattr(temp_stream.authenticator.active_token, 'github_organization', 'N/A')}")
         # replace manually provided org/repo values by the ones obtained
         # from github api. This guarantees that case is correct in the output data.
         # See https://github.com/MeltanoLabs/tap-github/issues/110
@@ -131,6 +147,7 @@ def validate_response(self, response: requests.Response) -> None:
                 if item == "rateLimit":
                     continue
                 try:
+                    self.logger.info(f"Processing record within TempStream: {record[item]}")
                     repo_full_name = "/".join(repo_list[int(item[4:])])
                     name_with_owner = record[item]["nameWithOwner"]
                     org, repo = name_with_owner.split("/")
@@ -172,17 +189,35 @@ def partitions(self) -> list[dict[str, str]] | None:
 
         if "repositories" in self.config:
             split_repo_names = [s.split("/") for s in self.config["repositories"]]
+
+            # Group repositories by organization for org-specific authentication
+            from collections import defaultdict
+            repos_by_org = defaultdict(list)
+            self.logger.info(f"Split repository names as part of repository stream: {split_repo_names}")
+            for org, repo in split_repo_names:
+                repos_by_org[org].append((org, repo))
+
             augmented_repo_list = []
             # chunk requests to the graphql endpoint to avoid timeouts and other
             # obscure errors that the api doesn't say much about. The actual limit
             # seems closer to 1000, use half that to stay safe.
             chunk_size = 500
             list_length = len(split_repo_names)
-            self.logger.info(f"Filtering repository list of {list_length} repositories")
-            for ndx in range(0, list_length, chunk_size):
-                augmented_repo_list += self.get_repo_ids(
-                    split_repo_names[ndx : ndx + chunk_size]
-                )
+            self.logger.info(f"Filtering repository list of {list_length} repositories across {len(repos_by_org)} organizations")
+
+            # Process each organization's repos separately with org-specific auth
+            for org, org_repos in repos_by_org.items():
+                self.logger.info(f"Validating {len(org_repos)} repositories for organization: {org}")
+                # Set organization-specific authentication before validating repos
+                self.logger.info(f"Setting organization within repository stream: {org}")
+                self.authenticator.set_organization(org)
+
+                # Process in chunks
+                for ndx in range(0, len(org_repos), chunk_size):
+                    augmented_repo_list += self.get_repo_ids(
+                        org_repos[ndx : ndx + chunk_size]
+                    )
+
             self.logger.info(
                 f"Running the tap on {len(augmented_repo_list)} repositories"
             )
@@ -216,6 +251,11 @@ def get_records(self, context: Context | None) -> Iterable[dict[str, Any]]:
         quota when only syncing a child stream. Without this,
         the API call is sent but data is discarded.
         """
+        # Set organization-specific authentication before fetching records
+        if context is not None and "org" in context:
+            self.logger.info(f"Setting organization within repository stream: {context['org']}")
+            self.authenticator.set_organization(context["org"])
+
         if (
             not self.selected
             and "skip_parent_streams" in self.config