Refactor query

Author: maxachis

src/api/endpoints/annotate/_shared/queries/helper.py

@@ -2,9 +2,12 @@
 This module contains helper functions for the annotate GET queries
 """
 
-from sqlalchemy import Select, case
+from sqlalchemy import Select, case, exists, select
 from sqlalchemy.orm import joinedload
 
+from src.collectors.enums import URLStatus
+from src.db.helpers.query import exists_url, not_exists_url
+from src.db.models.impl.flag.url_suspended.sqlalchemy import FlagURLSuspended
 from src.db.models.impl.url.core.enums import URLSource
 from src.db.models.impl.url.core.sqlalchemy import URL
 from src.db.models.views.unvalidated_url import UnvalidatedURL
@@ -15,11 +18,7 @@
 def get_select() -> Select:
     return (
         Select(URL)
-        # URL Must be unvalidated
-        .join(
-            UnvalidatedURL,
-            UnvalidatedURL.url_id == URL.id
-        )
+
         .join(
             URLAnnotationFlagsView,
             URLAnnotationFlagsView.url_id == URL.id
@@ -31,6 +30,19 @@ def get_select() -> Select:
     )
 
 def conclude(query: Select) -> Select:
+    # Add common where conditions
+    query = query.where(
+        URL.status == URLStatus.OK.value,
+        not_exists_url(
+            FlagURLSuspended
+        ),
+        # URL Must be unvalidated
+        exists_url(
+            UnvalidatedURL
+        )
+    )
+
+
     query = (
         # Add load options
         query.options(

src/api/endpoints/annotate/all/get/queries/core.py

@@ -4,17 +4,12 @@
 from src.api.endpoints.annotate._shared.extract import extract_and_format_get_annotation_result
 from src.api.endpoints.annotate._shared.queries import helper
 from src.api.endpoints.annotate.all.get.models.response import GetNextURLForAllAnnotationResponse
-from src.collectors.enums import URLStatus
 from src.db.models.impl.annotation.agency.user.sqlalchemy import AnnotationAgencyUser
 from src.db.models.impl.annotation.location.user.sqlalchemy import AnnotationLocationUser
-from src.db.models.impl.flag.url_suspended.sqlalchemy import FlagURLSuspended
-from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL
-from src.db.models.impl.url.core.sqlalchemy import URL
 from src.db.models.impl.annotation.record_type.user.user import AnnotationUserRecordType
 from src.db.models.impl.annotation.url_type.user.sqlalchemy import AnnotationUserURLType
-from src.db.models.views.unvalidated_url import UnvalidatedURL
-from src.db.models.views.url_anno_count import URLAnnotationCount
-from src.db.models.views.url_annotations_flags import URLAnnotationFlagsView
+from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL
+from src.db.models.impl.url.core.sqlalchemy import URL
 from src.db.queries.base.builder import QueryBuilderBase
 
 
@@ -45,7 +40,6 @@ async def run(
         query = (
             query
             .where(
-                    URL.status == URLStatus.OK.value,
                     # Must not have been previously annotated by user
                     ~exists(
                         select(AnnotationUserURLType.url_id)
@@ -78,14 +72,6 @@ async def run(
                             AnnotationUserRecordType.url_id == URL.id,
                             AnnotationUserRecordType.user_id == self.user_id,
                         )
-                    ),
-                    ~exists(
-                        select(
-                            FlagURLSuspended.url_id
-                        )
-                        .where(
-                            FlagURLSuspended.url_id == URL.id,
-                        )
                     )
             )
         )

src/api/endpoints/annotate/anonymous/get/query.py

@@ -1,27 +1,18 @@
-from typing import Any
 from uuid import UUID
 
-from sqlalchemy import Select, func, exists, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
 
 from src.api.endpoints.annotate._shared.extract import extract_and_format_get_annotation_result
+from src.api.endpoints.annotate._shared.queries import helper
 from src.api.endpoints.annotate.all.get.models.response import GetNextURLForAllAnnotationResponse
 from src.api.endpoints.annotate.anonymous.get.helpers import not_exists_anon_annotation
 from src.api.endpoints.annotate.anonymous.get.response import GetNextURLForAnonymousAnnotationResponse
-from src.collectors.enums import URLStatus
-from src.db.helpers.query import not_exists_url
 from src.db.models.impl.annotation.agency.anon.sqlalchemy import AnnotationAgencyAnon
 from src.db.models.impl.annotation.location.anon.sqlalchemy import AnnotationLocationAnon
 from src.db.models.impl.annotation.record_type.anon.sqlalchemy import AnnotationAnonRecordType
 from src.db.models.impl.annotation.url_type.anon.sqlalchemy import AnnotationAnonURLType
-from src.db.models.impl.flag.url_suspended.sqlalchemy import FlagURLSuspended
 from src.db.models.impl.url.core.sqlalchemy import URL
-from src.db.models.views.unvalidated_url import UnvalidatedURL
-from src.db.models.views.url_anno_count import URLAnnotationCount
-from src.db.models.views.url_annotations_flags import URLAnnotationFlagsView
 from src.db.queries.base.builder import QueryBuilderBase
-from src.api.endpoints.annotate._shared.queries import helper
 
 
 class GetNextURLForAnonymousAnnotationQueryBuilder(QueryBuilderBase):
@@ -40,7 +31,6 @@ async def run(self, session: AsyncSession) -> GetNextURLForAnonymousAnnotationRe
         query = (
             query
             .where(
-                URL.status == URLStatus.OK.value,
                 # Must not have been previously annotated by user
                 not_exists_anon_annotation(
                     session_id=self.session_id,
@@ -57,14 +47,6 @@ async def run(self, session: AsyncSession) -> GetNextURLForAnonymousAnnotationRe
                 not_exists_anon_annotation(
                     session_id=self.session_id,
                     anon_model=AnnotationAgencyAnon
-                ),
-                ~exists(
-                    select(
-                        FlagURLSuspended.url_id
-                    )
-                    .where(
-                        FlagURLSuspended.url_id == URL.id,
-                    )
                 )
             )
         )