🐛 Source Github: increase discovery time (#36429)

Author: davydov-d

airbyte-integrations/connectors/source-github/metadata.yaml

@@ -10,7 +10,7 @@ data:
   connectorSubtype: api
   connectorType: source
   definitionId: ef69ef6e-aa7f-4af1-a01d-ef775033524e
-  dockerImageTag: 1.7.0
+  dockerImageTag: 1.7.1
   dockerRepository: airbyte/source-github
   documentationUrl: https://docs.airbyte.com/integrations/sources/github
   githubIssueLabel: source-github

airbyte-integrations/connectors/source-github/pyproject.toml

@@ -3,7 +3,7 @@ requires = [ "poetry-core>=1.0.0",]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-version = "1.7.0"
+version = "1.7.1"
 name = "source-github"
 description = "Source implementation for GitHub."
 authors = [ "Airbyte <[email protected]>",]

airbyte-integrations/connectors/source-github/source_github/source.py

@@ -3,7 +3,7 @@
 #
 
 from os import getenv
-from typing import Any, Dict, List, Mapping, MutableMapping, Tuple
+from typing import Any, List, Mapping, MutableMapping, Optional, Tuple
 from urllib.parse import urlparse
 
 from airbyte_cdk import AirbyteLogger
@@ -65,7 +65,9 @@ class SourceGithub(AbstractSource):
     continue_sync_on_stream_failure = True
 
     @staticmethod
-    def _get_org_repositories(config: Mapping[str, Any], authenticator: MultipleTokenAuthenticator) -> Tuple[List[str], List[str]]:
+    def _get_org_repositories(
+        config: Mapping[str, Any], authenticator: MultipleTokenAuthenticator
+    ) -> Tuple[List[str], List[str], Optional[str]]:
         """
         Parse config/repositories and produce two lists: organizations, repositories.
         Args:
@@ -78,16 +80,19 @@ def _get_org_repositories(config: Mapping[str, Any], authenticator: MultipleToke
         organizations = set()
         unchecked_repos = set()
         unchecked_orgs = set()
+        pattern = None
 
         for org_repos in config_repositories:
-            org, _, repos = org_repos.partition("/")
-            if repos == "*":
-                unchecked_orgs.add(org)
+            _, _, repos = org_repos.partition("/")
+            if "*" in repos:
+                unchecked_orgs.add(org_repos)
             else:
                 unchecked_repos.add(org_repos)
 
         if unchecked_orgs:
-            stream = Repositories(authenticator=authenticator, organizations=unchecked_orgs, api_url=config.get("api_url"))
+            org_names = [org.split("/")[0] for org in unchecked_orgs]
+            pattern = "|".join([f"({org.replace('*', '.*')})" for org in unchecked_orgs])
+            stream = Repositories(authenticator=authenticator, organizations=org_names, api_url=config.get("api_url"), pattern=pattern)
             for record in read_full_refresh(stream):
                 repositories.add(record["full_name"])
                 organizations.add(record["organization"])
@@ -96,7 +101,7 @@ def _get_org_repositories(config: Mapping[str, Any], authenticator: MultipleToke
         if unchecked_repos:
             stream = RepositoryStats(
                 authenticator=authenticator,
-                repositories=unchecked_repos,
+                repositories=list(unchecked_repos),
                 api_url=config.get("api_url"),
                 # This parameter is deprecated and in future will be used sane default, page_size: 10
                 page_size_for_large_streams=config.get("page_size_for_large_streams", constants.DEFAULT_PAGE_SIZE_FOR_LARGE_STREAM),
@@ -107,7 +112,7 @@ def _get_org_repositories(config: Mapping[str, Any], authenticator: MultipleToke
                 if organization:
                     organizations.add(organization)
 
-        return list(organizations), list(repositories)
+        return list(organizations), list(repositories), pattern
 
     @staticmethod
     def get_access_token(config: Mapping[str, Any]):
@@ -169,45 +174,6 @@ def _validate_branches(self, config: MutableMapping[str, Any]) -> MutableMapping
     def _is_http_allowed() -> bool:
         return getenv("DEPLOYMENT_MODE", "").upper() != "CLOUD"
 
-    @staticmethod
-    def _get_branches_data(
-        selected_branches: List, full_refresh_args: Dict[str, Any] = None
-    ) -> Tuple[Dict[str, str], Dict[str, List[str]]]:
-        selected_branches = set(selected_branches)
-
-        # Get the default branch for each repository
-        default_branches = {}
-        repository_stats_stream = RepositoryStats(**full_refresh_args)
-        for stream_slice in repository_stats_stream.stream_slices(sync_mode=SyncMode.full_refresh):
-            default_branches.update(
-                {
-                    repo_stats["full_name"]: repo_stats["default_branch"]
-                    for repo_stats in repository_stats_stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
-                }
-            )
-
-        all_branches = []
-        branches_stream = Branches(**full_refresh_args)
-        for stream_slice in branches_stream.stream_slices(sync_mode=SyncMode.full_refresh):
-            for branch in branches_stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice):
-                all_branches.append(f"{branch['repository']}/{branch['name']}")
-
-        # Create mapping of repository to list of branches to pull commits for
-        # If no branches are specified for a repo, use its default branch
-        branches_to_pull: Dict[str, List[str]] = {}
-        for repo in full_refresh_args["repositories"]:
-            repo_branches = []
-            for branch in selected_branches:
-                branch_parts = branch.split("/", 2)
-                if "/".join(branch_parts[:2]) == repo and branch in all_branches:
-                    repo_branches.append(branch_parts[-1])
-            if not repo_branches:
-                repo_branches = [default_branches[repo]]
-
-            branches_to_pull[repo] = repo_branches
-
-        return default_branches, branches_to_pull
-
     def user_friendly_error_message(self, message: str) -> str:
         user_message = ""
         if "404 Client Error: Not Found for url: https://api.github.com/repos/" in message:
@@ -229,7 +195,7 @@ def check_connection(self, logger: AirbyteLogger, config: Mapping[str, Any]) ->
         config = self._validate_and_transform_config(config)
         try:
             authenticator = self._get_authenticator(config)
-            _, repositories = self._get_org_repositories(config=config, authenticator=authenticator)
+            _, repositories, _ = self._get_org_repositories(config=config, authenticator=authenticator)
             if not repositories:
                 return (
                     False,
@@ -246,7 +212,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         authenticator = self._get_authenticator(config)
         config = self._validate_and_transform_config(config)
         try:
-            organizations, repositories = self._get_org_repositories(config=config, authenticator=authenticator)
+            organizations, repositories, pattern = self._get_org_repositories(config=config, authenticator=authenticator)
         except Exception as e:
             message = repr(e)
             user_message = self.user_friendly_error_message(message)
@@ -291,7 +257,6 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         }
         repository_args_with_start_date = {**repository_args, "start_date": start_date}
 
-        default_branches, branches_to_pull = self._get_branches_data(config.get("branch", []), repository_args)
         pull_requests_stream = PullRequests(**repository_args_with_start_date)
         projects_stream = Projects(**repository_args_with_start_date)
         project_columns_stream = ProjectColumns(projects_stream, **repository_args_with_start_date)
@@ -307,7 +272,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
             Comments(**repository_args_with_start_date),
             CommitCommentReactions(**repository_args_with_start_date),
             CommitComments(**repository_args_with_start_date),
-            Commits(**repository_args_with_start_date, branches_to_pull=branches_to_pull, default_branches=default_branches),
+            Commits(**repository_args_with_start_date, branches_to_pull=config.get("branches", [])),
             ContributorActivity(**repository_args),
             Deployments(**repository_args_with_start_date),
             Events(**repository_args_with_start_date),
@@ -327,7 +292,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
             ProjectsV2(**repository_args_with_start_date),
             pull_requests_stream,
             Releases(**repository_args_with_start_date),
-            Repositories(**organization_args_with_start_date),
+            Repositories(**organization_args_with_start_date, pattern=pattern),
             ReviewComments(**repository_args_with_start_date),
             Reviews(**repository_args_with_start_date),
             Stargazers(**repository_args_with_start_date),

airbyte-integrations/connectors/source-github/source_github/spec.json

@@ -81,18 +81,19 @@
         "type": "array",
         "items": {
           "type": "string",
-          "pattern": "^([\\w.-]+/(\\*|[\\w.-]+(?<!\\.git))\\s+)*[\\w.-]+/(\\*|[\\w.-]+(?<!\\.git))$"
+          "pattern": "^[\\w.-]+/(([\\w.-]*\\*)|[\\w.-]+(?<!\\.git))$"
         },
         "minItems": 1,
         "examples": [
-          "airbytehq/airbyte airbytehq/another-repo",
+          "airbytehq/airbyte",
+          "airbytehq/another-repo",
           "airbytehq/*",
-          "airbytehq/airbyte"
+          "airbytehq/a*"
         ],
         "title": "GitHub Repositories",
-        "description": "List of GitHub organizations/repositories, e.g. `airbytehq/airbyte` for single repository, `airbytehq/*` for get all repositories from organization and `airbytehq/airbyte airbytehq/another-repo` for multiple repositories.",
+        "description": "List of GitHub organizations/repositories, e.g. `airbytehq/airbyte` for single repository, `airbytehq/*` for get all repositories from organization and `airbytehq/a* for matching multiple repositories by pattern.",
         "order": 1,
-        "pattern_descriptor": "org/repo org/another-repo org/*"
+        "pattern_descriptor": "org/repo org/another-repo org/* org/a*"
       },
       "start_date": {
         "type": "string",
@@ -126,7 +127,7 @@
           "type": "string"
         },
         "title": "Branches",
-        "examples": ["airbytehq/airbyte/master airbytehq/airbyte/my-branch"],
+        "examples": ["airbytehq/airbyte/master", "airbytehq/airbyte/my-branch"],
         "description": "List of GitHub repository branches to pull commits for, e.g. `airbytehq/airbyte/master`. If no branches are specified for a repository, the default branch will be pulled.",
         "order": 4,
         "pattern_descriptor": "org/repo/branch1 org/repo/branch2"

airbyte-integrations/connectors/source-github/source_github/streams.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+import re
 import time
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional
@@ -441,12 +442,18 @@ class Repositories(SemiIncrementalMixin, Organizations):
         "direction": "desc",
     }
 
+    def __init__(self, *args, pattern: Optional[str] = None, **kwargs):
+        self._pattern = re.compile(pattern) if pattern else pattern
+        super().__init__(*args, **kwargs)
+
     def path(self, stream_slice: Mapping[str, Any] = None, **kwargs) -> str:
         return f"orgs/{stream_slice['organization']}/repos"
 
     def parse_response(self, response: requests.Response, stream_slice: Mapping[str, Any] = None, **kwargs) -> Iterable[Mapping]:
         for record in response.json():  # GitHub puts records in an array.
-            yield self.transform(record=record, stream_slice=stream_slice)
+            record = self.transform(record=record, stream_slice=stream_slice)
+            if not self._pattern or self._pattern.match(record["full_name"]):
+                yield record
 
 
 class Tags(GithubStream):
@@ -676,10 +683,13 @@ class Commits(IncrementalMixin, GithubStream):
     cursor_field = "created_at"
     slice_keys = ["repository", "branch"]
 
-    def __init__(self, branches_to_pull: Mapping[str, List[str]], default_branches: Mapping[str, str], **kwargs):
+    def __init__(self, branches_to_pull: List[str], **kwargs):
         super().__init__(**kwargs)
-        self.branches_to_pull = branches_to_pull
-        self.default_branches = default_branches
+        kwargs.pop("start_date")
+        self.branches_to_repos = {}
+        self.branches_to_pull = set(branches_to_pull)
+        self.branches_stream = Branches(**kwargs)
+        self.repositories_stream = RepositoryStats(**kwargs)
 
     def request_params(self, stream_state: Mapping[str, Any], stream_slice: Mapping[str, Any] = None, **kwargs) -> MutableMapping[str, Any]:
         params = super(IncrementalMixin, self).request_params(stream_state=stream_state, stream_slice=stream_slice, **kwargs)
@@ -690,9 +700,10 @@ def request_params(self, stream_state: Mapping[str, Any], stream_slice: Mapping[
         return params
 
     def stream_slices(self, **kwargs) -> Iterable[Optional[Mapping[str, Any]]]:
+        self._validate_branches_to_pull()
         for stream_slice in super().stream_slices(**kwargs):
             repository = stream_slice["repository"]
-            for branch in self.branches_to_pull.get(repository, []):
+            for branch in self.branches_to_repos.get(repository, []):
                 yield {"branch": branch, "repository": repository}
 
     def transform(self, record: MutableMapping[str, Any], stream_slice: Mapping[str, Any]) -> MutableMapping[str, Any]:
@@ -718,6 +729,30 @@ def get_updated_state(self, current_stream_state: MutableMapping[str, Any], late
         current_stream_state.setdefault(repository, {}).setdefault(branch, {})[self.cursor_field] = updated_state
         return current_stream_state
 
+    def _validate_branches_to_pull(self):
+        # Get the default branch for each repository
+        default_branches = {}
+        for stream_slice in self.repositories_stream.stream_slices(sync_mode=SyncMode.full_refresh):
+            for repo_stats in self.repositories_stream.read_records(stream_slice=stream_slice, sync_mode=SyncMode.full_refresh):
+                default_branches[repo_stats["full_name"]] = repo_stats["default_branch"]
+
+        all_branches = []
+        for stream_slice in self.branches_stream.stream_slices(sync_mode=SyncMode.full_refresh):
+            for branch in self.branches_stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice):
+                all_branches.append(f"{branch['repository']}/{branch['name']}")
+
+        # Create mapping of repository to list of branches to pull commits for
+        # If no branches are specified for a repo, use its default branch
+        for repo in self.repositories:
+            repo_branches = []
+            for branch in self.branches_to_pull:
+                branch_parts = branch.split("/", 2)
+                if "/".join(branch_parts[:2]) == repo and branch in all_branches:
+                    repo_branches.append(branch_parts[-1])
+            if not repo_branches:
+                repo_branches = [default_branches[repo]]
+            self.branches_to_repos[repo] = repo_branches
+
 
 class Issues(IncrementalMixin, GithubStream):
     """

airbyte-integrations/connectors/source-github/unit_tests/test_source.py

@@ -105,7 +105,7 @@ def test_check_connection_repos_and_org_repos(rate_limit_mock_response):
 
 @responses.activate
 def test_check_connection_org_only(rate_limit_mock_response):
-    repos = [{"name": f"name {i}", "full_name": f"full name {i}", "updated_at": "2020-01-01T00:00:00Z"} for i in range(1000)]
+    repos = [{"name": f"name {i}", "full_name": f"airbytehq/full name {i}", "updated_at": "2020-01-01T00:00:00Z"} for i in range(1000)]
     responses.add("GET", "https://api.github.com/orgs/airbytehq/repos", json=repos)
 
     status = check_source("airbytehq/*")
@@ -115,49 +115,6 @@ def test_check_connection_org_only(rate_limit_mock_response):
     assert len(responses.calls) == 2
 
 
-@responses.activate
-def test_get_branches_data():
-
-    repository_args = {"repositories": ["airbytehq/integration-test"], "page_size_for_large_streams": 10}
-
-    source = SourceGithub()
-
-    responses.add(
-        "GET",
-        "https://api.github.com/repos/airbytehq/integration-test",
-        json={"full_name": "airbytehq/integration-test", "default_branch": "master"},
-    )
-
-    responses.add(
-        "GET",
-        "https://api.github.com/repos/airbytehq/integration-test/branches",
-        json=[
-            {"repository": "airbytehq/integration-test", "name": "feature/branch_0"},
-            {"repository": "airbytehq/integration-test", "name": "feature/branch_1"},
-            {"repository": "airbytehq/integration-test", "name": "feature/branch_2"},
-            {"repository": "airbytehq/integration-test", "name": "master"},
-        ],
-    )
-
-    default_branches, branches_to_pull = source._get_branches_data([], repository_args)
-    assert default_branches == {"airbytehq/integration-test": "master"}
-    assert branches_to_pull == {"airbytehq/integration-test": ["master"]}
-
-    default_branches, branches_to_pull = source._get_branches_data(
-        [
-            "airbytehq/integration-test/feature/branch_0",
-            "airbytehq/integration-test/feature/branch_1",
-            "airbytehq/integration-test/feature/branch_3",
-        ],
-        repository_args,
-    )
-
-    assert default_branches == {"airbytehq/integration-test": "master"}
-    assert len(branches_to_pull["airbytehq/integration-test"]) == 2
-    assert "feature/branch_0" in branches_to_pull["airbytehq/integration-test"]
-    assert "feature/branch_1" in branches_to_pull["airbytehq/integration-test"]
-
-
 @responses.activate
 def test_get_org_repositories():
     responses.add(
@@ -178,15 +135,15 @@ def test_get_org_repositories():
     config = {"repositories": ["airbytehq/integration-test", "docker/*"]}
     source = SourceGithub()
     config = source._ensure_default_values(config)
-    organisations, repositories = source._get_org_repositories(config, authenticator=None)
+    organisations, repositories, _ = source._get_org_repositories(config, authenticator=None)
 
     assert set(repositories) == {"airbytehq/integration-test", "docker/docker-py", "docker/compose"}
     assert set(organisations) == {"airbytehq", "docker"}
 
 
 @responses.activate
 def test_organization_or_repo_available(monkeypatch, rate_limit_mock_response):
-    monkeypatch.setattr(SourceGithub, "_get_org_repositories", MagicMock(return_value=(False, False)))
+    monkeypatch.setattr(SourceGithub, "_get_org_repositories", MagicMock(return_value=(False, False, None)))
     source = SourceGithub()
     with pytest.raises(Exception) as exc_info:
         config = {"access_token": "test_token", "repository": ""}
@@ -209,6 +166,7 @@ def test_check_config_repository():
         "airbyte_hq/airbyte",
         "airbytehq/123",
         "airbytehq/airbytexgit",
+        "airbytehq/a*",
     ]
 
     repos_fail = [
@@ -242,7 +200,7 @@ def test_check_config_repository():
 
 @responses.activate
 def test_streams_no_streams_available_error(monkeypatch, rate_limit_mock_response):
-    monkeypatch.setattr(SourceGithub, "_get_org_repositories", MagicMock(return_value=(False, False)))
+    monkeypatch.setattr(SourceGithub, "_get_org_repositories", MagicMock(return_value=(False, False, None)))
     with pytest.raises(AirbyteTracedException) as e:
         SourceGithub().streams(config={"access_token": "test_token", "repository": "airbytehq/airbyte-test"})
     assert str(e.value) == "No streams available. Please check permissions"