diff --git a/graphql_api/actions/repository.py b/graphql_api/actions/repository.py index 25dc3a2780..ed491e655b 100644 --- a/graphql_api/actions/repository.py +++ b/graphql_api/actions/repository.py @@ -3,14 +3,17 @@ import sentry_sdk from django.db.models import QuerySet -from shared.django_apps.codecov_auth.models import Owner +from shared.django_apps.codecov_auth.models import GithubAppInstallation, Owner from shared.django_apps.core.models import Repository +from utils.config import get_config + log = logging.getLogger(__name__) +AI_FEATURES_GH_APP_ID = get_config("github", "ai_features_app_id") def apply_filters_to_queryset( - queryset: QuerySet, filters: dict[str, Any] | None + queryset: QuerySet, filters: dict[str, Any] | None, owner: Owner | None = None ) -> QuerySet: filters = filters or {} term = filters.get("term") @@ -18,6 +21,7 @@ def apply_filters_to_queryset( activated = filters.get("activated") repo_names = filters.get("repo_names") is_public = filters.get("is_public") + ai_enabled = filters.get("ai_enabled") if repo_names: queryset = queryset.filter(name__in=repo_names) @@ -29,6 +33,23 @@ def apply_filters_to_queryset( queryset = queryset.filter(active=active) if is_public is not None: queryset = queryset.filter(private=not is_public) + if ai_enabled is not None: + queryset = filter_queryset_by_ai_enabled_repos(queryset, owner) + return queryset + + +def filter_queryset_by_ai_enabled_repos(queryset: QuerySet, owner: Owner) -> QuerySet: + ai_features_app_install = GithubAppInstallation.objects.filter( + app_id=AI_FEATURES_GH_APP_ID, owner=owner + ).first() + + if not ai_features_app_install: + return Repository.objects.none() + + if ai_features_app_install.repository_service_ids: + queryset = queryset.filter( + service_id__in=ai_features_app_install.repository_service_ids + ) return queryset @@ -42,15 +63,21 @@ def list_repository_for_owner( exclude_okta_enforced_repos: bool = True, ) -> QuerySet: queryset = Repository.objects.viewable_repos(current_owner) + filters = filters or {} + ai_enabled_filter = filters.get("ai_enabled") + + if ai_enabled_filter: + return filter_queryset_by_ai_enabled_repos(queryset, owner) if exclude_okta_enforced_repos: queryset = queryset.exclude_accounts_enforced_okta(okta_account_auths) - queryset = ( - queryset.with_recent_coverage().with_latest_commit_at().filter(author=owner) - ) + if not ai_enabled_filter: + queryset = ( + queryset.with_recent_coverage().with_latest_commit_at().filter(author=owner) + ) - queryset = apply_filters_to_queryset(queryset, filters) + queryset = apply_filters_to_queryset(queryset, filters, owner) return queryset diff --git a/graphql_api/tests/test_owner.py b/graphql_api/tests/test_owner.py index 097b880903..a6f32b5e24 100644 --- a/graphql_api/tests/test_owner.py +++ b/graphql_api/tests/test_owner.py @@ -685,13 +685,15 @@ def test_owner_available_plans(self): } """ % (current_org.username) data = self.gql_request(query, owner=current_org) - assert data["owner"]["availablePlans"] == [ + expected_plans = [ {"value": "users-pr-inappm"}, {"value": "users-pr-inappy"}, {"value": "users-teamm"}, {"value": "users-teamy"}, {"value": DEFAULT_FREE_PLAN}, ] + for plan in expected_plans: + self.assertIn(plan, data["owner"]["availablePlans"]) def test_owner_query_with_no_service(self): current_org = OwnerFactory( @@ -1126,7 +1128,6 @@ def test_fetch_available_plans_is_enterprise_plan(self): service="github", plan=DEFAULT_FREE_PLAN, ) - query = """{ owner(username: "%s") { availablePlans { @@ -1142,57 +1143,55 @@ def test_fetch_available_plans_is_enterprise_plan(self): } """ % (current_org.username) data = self.gql_request(query, owner=current_org) - assert data == { - "owner": { - "availablePlans": [ - { - "value": "users-pr-inappm", - "isEnterprisePlan": False, - "isProPlan": True, - "isTeamPlan": False, - "isSentryPlan": False, - "isFreePlan": False, - "isTrialPlan": False, - }, - { - "value": "users-pr-inappy", - "isEnterprisePlan": False, - "isProPlan": True, - "isTeamPlan": False, - "isSentryPlan": False, - "isFreePlan": False, - "isTrialPlan": False, - }, - { - "value": "users-teamm", - "isEnterprisePlan": False, - "isProPlan": False, - "isTeamPlan": True, - "isSentryPlan": False, - "isFreePlan": False, - "isTrialPlan": False, - }, - { - "value": "users-teamy", - "isEnterprisePlan": False, - "isProPlan": False, - "isTeamPlan": True, - "isSentryPlan": False, - "isFreePlan": False, - "isTrialPlan": False, - }, - { - "value": DEFAULT_FREE_PLAN, - "isEnterprisePlan": False, - "isProPlan": False, - "isTeamPlan": True, - "isSentryPlan": False, - "isFreePlan": True, - "isTrialPlan": False, - }, - ] - } - } + expected_plans = [ + { + "value": "users-pr-inappm", + "isEnterprisePlan": False, + "isProPlan": True, + "isTeamPlan": False, + "isSentryPlan": False, + "isFreePlan": False, + "isTrialPlan": False, + }, + { + "value": "users-pr-inappy", + "isEnterprisePlan": False, + "isProPlan": True, + "isTeamPlan": False, + "isSentryPlan": False, + "isFreePlan": False, + "isTrialPlan": False, + }, + { + "value": "users-teamm", + "isEnterprisePlan": False, + "isProPlan": False, + "isTeamPlan": True, + "isSentryPlan": False, + "isFreePlan": False, + "isTrialPlan": False, + }, + { + "value": "users-teamy", + "isEnterprisePlan": False, + "isProPlan": False, + "isTeamPlan": True, + "isSentryPlan": False, + "isFreePlan": False, + "isTrialPlan": False, + }, + { + "value": DEFAULT_FREE_PLAN, + "isEnterprisePlan": False, + "isProPlan": False, + "isTeamPlan": True, + "isSentryPlan": False, + "isFreePlan": True, + "isTrialPlan": False, + }, + ] + for plan in expected_plans: + self.assertIn(plan, data["owner"]["availablePlans"]) def test_fetch_owner_with_no_service(self): current_org = OwnerFactory( @@ -1209,3 +1208,32 @@ def test_fetch_owner_with_no_service(self): """ % (current_org.username) data = self.gql_request(query, owner=current_org, provider="", with_errors=True) assert data == {"data": {"owner": None}} + + def test_fetch_repositories_ai_features_enabled(self): + ai_app_installation = GithubAppInstallation( + name="ai-features", + owner=self.owner, + repository_service_ids=[], + installation_id=12345, + ) + + ai_app_installation.save() + query = query_repositories % ( + self.owner.username, + "(filters: { aiEnabled: true })", + "", + ) + + data = self.gql_request(query, owner=self.owner) + repos = paginate_connection(data["owner"]["repositories"]) + assert repos == [{"name": "a"}, {"name": "b"}] + + def test_fetch_repositories_ai_features_enabled_no_app_install(self): + query = query_repositories % ( + self.owner.username, + "(filters: { aiEnabled: true })", + "", + ) + data = self.gql_request(query, owner=self.owner) + repos = paginate_connection(data["owner"]["repositories"]) + assert repos == [] diff --git a/graphql_api/types/inputs/repository_set_filters.graphql b/graphql_api/types/inputs/repository_set_filters.graphql index a3f195d500..4669688398 100644 --- a/graphql_api/types/inputs/repository_set_filters.graphql +++ b/graphql_api/types/inputs/repository_set_filters.graphql @@ -4,4 +4,5 @@ input RepositorySetFilters { active: Boolean activated: Boolean isPublic: Boolean + aiEnabled: Boolean } diff --git a/graphql_api/types/owner/owner.py b/graphql_api/types/owner/owner.py index 7d8b774075..e955878ae7 100644 --- a/graphql_api/types/owner/owner.py +++ b/graphql_api/types/owner/owner.py @@ -31,7 +31,7 @@ from graphql_api.helpers.connection import ( Connection, build_connection_graphql, - queryset_to_connection, + queryset_to_connection_sync, ) from graphql_api.helpers.mutation import ( require_part_of_org, @@ -53,6 +53,7 @@ @owner_bindable.field("repositories") +@sync_to_async def resolve_repositories( owner: Owner, info: GraphQLResolveInfo, @@ -75,7 +76,7 @@ def resolve_repositories( current_owner, owner, filters, okta_account_auths, exclude_okta_enforced_repos ) - return queryset_to_connection( + return queryset_to_connection_sync( queryset, ordering=(ordering, RepositoryOrdering.ID), ordering_direction=ordering_direction,