[feat] Enable pagination for AI enabled repos

graphql_api/actions/repository.py

@@ -3,21 +3,25 @@
 
 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")
     active = filters.get("active")
     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 @@
         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 @@
     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
 
 

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 == []

graphql_api/types/inputs/repository_set_filters.graphql

@@ -4,4 +4,5 @@ input RepositorySetFilters {
   active: Boolean
   activated: Boolean
   isPublic: Boolean
+  aiEnabled: Boolean
 }

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,