Repository: oppia/oppia
File: core/jobs/batch_jobs/audit_threads_with_missing_suggestions_jobs.py (Python)

Code before:
"""Audit job that finds feedback threads marked as having suggestions
but without a corresponding GeneralSuggestionModel.
"""

from __future__ import annotations

from core.jobs import base_jobs
from core.jobs.io import ndb_io
from core.jobs.types import job_run_result
from core.platform import models

import apache_beam as beam

MYPY = False
if MYPY:  # pragma: no cover
    from mypy_imports import feedback_models, suggestion_models

(feedback_models, suggestion_models) = models.Registry.import_models(
    [
        models.Names.FEEDBACK,
        models.Names.SUGGESTION,
    ]
)


class BaseThreadsWithMissingSuggestionsJob(base_jobs.JobBase):
    """Base class for threads with has_suggestion=True but no suggestion."""

    DATASTORE_UPDATES_ALLOWED = False

    def _get_suggestion_thread_ids(self) -> beam.PCollection[str]:
        """Returns thread_ids that have a GeneralSuggestionModel."""
        return (
            self.pipeline
            | 'Get GeneralSuggestionModels'
            >> ndb_io.GetModels(
                suggestion_models.GeneralSuggestionModel.get_all(
                    include_deleted=False
                )
            )
            | 'Extract suggestion thread ids'
            >> beam.Map(lambda model: model.id)
        )

    def _get_invalid_threads(
        self,
        valid_suggestion_thread_ids: beam.PCollection[str],
    ) -> beam.PCollection[feedback_models.GeneralFeedbackThreadModel]:
        """Returns threads marked as having suggestions but without one."""
        return (
            self.pipeline

Review comment (suggestion):
GCP will lose too much information by using private methods, we should prefer PTransforms as much as possible to make these re-usable and easier to inspect for release coordinators:

```python
class _GetSuggestionThreadIds(beam.PTransform):
    def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection[str]:
        # identical to `_get_suggestion_thread_ids`, but use `pbegin` rather
        # than `self.pipeline` to keep the job's "graph" accurate in GCP.


class _GetInvalidFeedbackThreads(beam.PTransform):
    def expand(
        self, valid_suggestion_thread_ids: beam.PCollection[str]
    ) -> beam.PCollection[feedback_models.GeneralFeedbackThreadModel]:
        # As you have it OR take a look at my next (non-blocking) comment.
```

Your tests look good, so I don't think you'll have to rewrite them by using this approach.

Code after:
"""Audit job that finds feedback threads marked as having suggestions
but without a corresponding GeneralSuggestionModel.
"""

from __future__ import annotations

from core.jobs import base_jobs
from core.jobs.io import ndb_io
from core.jobs.types import job_run_result
from core.platform import models

import apache_beam as beam

MYPY = False
if MYPY:  # pragma: no cover
    from mypy_imports import feedback_models, suggestion_models

(feedback_models, suggestion_models) = models.Registry.import_models(
    [
        models.Names.FEEDBACK,
        models.Names.SUGGESTION,
    ]
)


class FixThreadsWithMissingSuggestionsJob(base_jobs.JobBase):
    """Fixes feedback threads where has_suggestion=True but no suggestion exists.

    When DATASTORE_UPDATES_ALLOWED is False, this job behaves as an audit job
    and only reports invalid models without mutating the datastore.
    """

    DATASTORE_UPDATES_ALLOWED = True

    def run(self) -> beam.PCollection[job_run_result.JobRunResult]:
        """Runs the job.

        Returns:
            PCollection[JobRunResult]. Audit results when datastore updates
            are disabled, otherwise fix results.
        """

        suggestion_thread_ids = (
            self.pipeline
            | 'Get GeneralSuggestionModels'
            >> ndb_io.GetModels(
                suggestion_models.GeneralSuggestionModel.get_all(
                    include_deleted=False
                )
            )
            | 'Extract suggestion thread ids'