Try new approach

Author: rohitvinnakota-codecov

graphql_api/actions/repository.py

@@ -3,21 +3,24 @@
 
 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 graphql_api.types.owner.owner import AI_FEATURES_GH_APP_ID
+
 log = logging.getLogger(__name__)
 
 
 def apply_filters_to_queryset(
-    queryset: QuerySet, filters: dict[str, Any] | None
+    queryset: QuerySet, filters: dict[str, Any] | None, owner: Owner
 ) -> 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 +32,25 @@ def apply_filters_to_queryset(
         queryset = queryset.filter(active=active)
     if is_public is not None:
         queryset = queryset.filter(private=not is_public)
+    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
 
@@ -43,14 +65,20 @@ def list_repository_for_owner(
 ) -> QuerySet:
     queryset = Repository.objects.viewable_repos(current_owner)
 
+    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

@@ -1194,65 +1194,6 @@ def test_fetch_available_plans_is_enterprise_plan(self):
             }
         }
 
-    @patch("services.self_hosted.get_config")
-    def test_ai_enabled_repositories(self, get_config_mock):
-        get_config_mock.return_value = [
-            {"service": "github", "ai_features_app_id": 12345},
-        ]
-
-        ai_app_installation = GithubAppInstallation(
-            name="ai-features",
-            owner=self.owner,
-            repository_service_ids=[],
-            installation_id=12345,
-        )
-
-        ai_app_installation.save()
-
-        query = """{
-            owner(username: "%s") {
-                aiEnabledRepositories(first: 20) {
-                    edges {
-                        node {
-                            name
-                        }
-                    }
-                }
-            }
-        }
-
-        """ % (self.owner.username)
-        data = self.gql_request(query, owner=self.owner)
-        repos = paginate_connection(data["owner"]["aiEnabledRepositories"])
-        assert repos == [{"name": "a"}, {"name": "b"}]
-
-    @patch("services.self_hosted.get_config")
-    def test_ai_enabled_repositories_app_not_configured(self, get_config_mock):
-        current_org = OwnerFactory(
-            username="random-plan-user",
-            service="github",
-        )
-
-        get_config_mock.return_value = [
-            {"service": "github", "ai_features_app_id": 12345},
-        ]
-
-        query = """{
-            owner(username: "%s") {
-                aiEnabledRepositories {
-                    edges {
-                        node {
-                            name
-                        }
-                    }
-                }
-            }
-        }
-
-        """ % (current_org.username)
-        data = self.gql_request(query, owner=current_org)
-        assert data["owner"]["aiEnabledRepositories"] is None
-
     def test_fetch_owner_with_no_service(self):
         current_org = OwnerFactory(
             username="random-plan-user",

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.graphql

@@ -40,14 +40,6 @@ type Owner {
   yaml: String
   aiFeaturesEnabled: Boolean!
   aiEnabledRepos: [String]
-  aiEnabledRepositories(
-    ordering: RepositoryOrdering
-    orderingDirection: OrderingDirection
-    first: Int
-    after: String
-    last: Int
-    before: String
-  ): RepositoryConnection @cost(complexity: 25, multipliers: ["first", "last"])
   uploadTokenRequired: Boolean
   activatedUserCount: Int
 }

graphql_api/types/owner/owner.py

@@ -31,7 +31,6 @@
 from graphql_api.helpers.connection import (
     Connection,
     build_connection_graphql,
-    queryset_to_connection,
     queryset_to_connection_sync,
 )
 from graphql_api.helpers.mutation import (
@@ -55,6 +54,7 @@
 
 
 @owner_bindable.field("repositories")
+@sync_to_async
 def resolve_repositories(
     owner: Owner,
     info: GraphQLResolveInfo,
@@ -77,7 +77,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,
@@ -406,34 +406,3 @@ def resolve_activated_user_count(owner: Owner, info: GraphQLResolveInfo) -> int:
 @require_part_of_org
 def resolve_billing(owner: Owner, info: GraphQLResolveInfo) -> dict | None:
     return owner
-
-
-@owner_bindable.field("aiEnabledRepositories")
-@sync_to_async
-def resolve_ai_enabled_repositories(
-    owner: Owner,
-    info: GraphQLResolveInfo,
-    ordering: Optional[RepositoryOrdering] = RepositoryOrdering.ID,
-    ordering_direction: Optional[OrderingDirection] = OrderingDirection.ASC,
-    **kwargs: Any,
-) -> Coroutine[Any, Any, Connection] | None:
-    ai_features_app_install = GithubAppInstallation.objects.filter(
-        app_id=AI_FEATURES_GH_APP_ID, owner=owner
-    ).first()
-
-    if not ai_features_app_install:
-        return None
-
-    current_owner = info.context["request"].current_owner
-    queryset = Repository.objects.filter(author=owner).viewable_repos(current_owner)
-    if ai_features_app_install.repository_service_ids:
-        queryset = queryset.filter(
-            service_id__in=ai_features_app_install.repository_service_ids
-        )
-
-    return queryset_to_connection_sync(
-        queryset,
-        ordering=(ordering, RepositoryOrdering.ID),
-        ordering_direction=ordering_direction,
-        **kwargs,
-    )