53a4fc85-bb39-4dee-81e5-c2abd320ed53

Committed 15 Jul 2024 06:57PM UTC coverage: 85.407% (-0.01%) from 85.42%

Build # 53a4fc85-bb39-4dee-81e5-c2abd320ed53

Build Type

Pull #4880

circleci

Committed by

jwhitlock

Commit Message

Run each cleaner task in transaction

Pull Request Pull Request #4880: MPP-3846: Allow percent greater than 100%, and run in transaction

Run Details

4086 of 5235 branches covered (78.05%)

Branch coverage included in aggregate %.

26 of 27 new or added lines in 2 files covered. (96.3%)

1 existing line in 1 file now uncovered.

15913 of 18181 relevant lines covered (87.53%)

10.98 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.74

/privaterelay/management/commands/cleanup_data.py

from __future__ import annotations

import logging
import textwrap
from argparse import RawDescriptionHelpFormatter
from shutil import get_terminal_size
from typing import TYPE_CHECKING, Any

from django.core.management.base import BaseCommand, DjangoHelpFormatter
from django.db import transaction

from codetiming import Timer

from emails.cleaners import MissingProfileCleaner, ServerStorageCleaner
from privaterelay.cleaners import MissingEmailCleaner

if TYPE_CHECKING:  # pragma: no cover
    from argparse import ArgumentParser

    from privaterelay.cleaner_task import DataIssueTask


logger = logging.getLogger("eventsinfo.cleanup_data")


class RawDescriptionDjangoHelpFormatter(
    DjangoHelpFormatter, RawDescriptionHelpFormatter
):
    """DjangoHelpFormatter, but don't reflow the epilog."""


class Command(BaseCommand):
    help = "Detects and optionally clean data issues."
    task_list: list[type[DataIssueTask]] = [
        ServerStorageCleaner,
        MissingProfileCleaner,
        MissingEmailCleaner,
    ]
    tasks: dict[str, DataIssueTask]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        tasks = {Task.slug: Task() for Task in self.task_list}
        self.tasks = {slug: tasks[slug] for slug in sorted(tasks.keys())}

    def create_parser(self, prog_name, subcommand, **kwargs):
        """Add DataIssueTask information to default parser."""

        epilog_lines = ["Data issue detectors / cleaners:"]
        for slug, task in self.tasks.items():
            raw = f"{slug} - {task.title}"
            epilog_lines.extend(
                textwrap.wrap(
                    raw,
                    width=get_terminal_size().columns,
                    initial_indent="  ",
                    subsequent_indent="      ",
                )
            )
        epilog = "\n".join(epilog_lines)

        parser = super().create_parser(prog_name, subcommand, epilog=epilog, **kwargs)
        if not parser.formatter_class == DjangoHelpFormatter:
            raise TypeError("parser.formatter_class must be type DjangoHelpFormatter")
        parser.formatter_class = RawDescriptionDjangoHelpFormatter
        return parser

    def add_arguments(self, parser: ArgumentParser) -> None:
        parser.add_argument(
            "--clean", action="store_true", help="Clean detected data issues."
        )
        for slug, task in self.tasks.items():
            parser.add_argument(
                f"--{slug}",
                action="store_true",
                help=f"Only run {slug} {'cleaner' if task.can_clean else 'detector'}",
            )

    def handle(self, clean: bool, verbosity: int, *args: Any, **kwargs: Any) -> str:
        """Run the cleanup_data command."""

        # Determine if we're running some tasks or the full set
        run_some: list[str] = []
        for name, val in kwargs.items():
            slug = name.replace("_", "-")
            if slug in self.tasks and val:
                run_some.append(slug)
        if run_some:
            tasks = {slug: self.tasks[slug] for slug in sorted(run_some)}
        else:
            tasks = self.tasks

        # Find data issues and clean them if requested
        issue_count = 0
        clean_count = 0
        issue_timers: dict[str, float] = {}
        clean_timers: dict[str, float] = {}
        for slug, task in tasks.items():
            with transaction.atomic():
                task_issues, task_time = self.find_issues_in_task(task)
                issue_count += task_issues
                issue_timers[slug] = task_time

                if clean:
                    task_cleaned, clean_time = self.clean_issues_in_task(task)
                    if task_cleaned is not None and clean_time is not None:
                        clean_count += task_cleaned
                        clean_timers[slug] = clean_time

        # Log results and create a report, based on requested verbosity
        if clean:
            log_message = (
                f"cleanup_data complete, cleaned {clean_count} of"
                f" {issue_count} issue{'' if issue_count==1 else 's'}."
            )
        else:
            log_message = (
                f"cleanup_data complete, found {issue_count}"
                f" issue{'' if issue_count==1 else 's'} (dry run)."
            )

        full_data, log_data = self.prepare_data(
            clean, verbosity, tasks, issue_timers, clean_timers=clean_timers
        )
        logger.info(log_message, extra=log_data)
        report = self.get_report(log_message, full_data)
        return report

    def find_issues_in_task(self, task: DataIssueTask) -> tuple[int, float]:
        """Run a task, returning the number of issues and execution time."""
        with Timer(logger=None) as issue_timer:
            total = task.issues()
        return total, round(issue_timer.last, 3)

    def clean_issues_in_task(
        self, task: DataIssueTask
    ) -> tuple[int, float] | tuple[None, None]:
        """Run the task cleaner, returning issues cleaned and execution time."""
        if not hasattr(task, "clean"):
            return (None, None)
        with Timer(logger=None) as clean_timer:
            cleaned = task.clean()
        return cleaned, round(clean_timer.last, 3)

    def prepare_data(
        self,
        cleaned: bool,
        verbosity: int,
        tasks: dict[str, DataIssueTask],
        issue_timers: dict[str, float],
        clean_timers: dict[str, float],
    ) -> tuple[dict, dict]:
        """Gather full data and log data from tasks and timers."""

        timers = {"query_s": sum(issue_timers.values())}
        if clean_timers:
            timers["clean_s"] = sum(clean_timers.values())

        log_details = {}
        full_details = {}
        for slug, task in tasks.items():
            timers = {"query_s": issue_timers[slug]}
            if clean_timers:
                timers["clean_s"] = clean_timers[slug]

            full_details[slug] = {
                "title": task.title,
                "check_description": task.check_description,
                "can_clean": task.can_clean,
                "counts": task.counts,
                "markdown_report": task.markdown_report(),
                "timers": timers,
            }

            if verbosity > 1:
                log_details[slug] = {"counts": task.counts, "timers": timers}
            elif verbosity > 0:
                log_details[slug] = {"counts": {"summary": task.counts["summary"]}}

        full_data = {
            "verbosity": verbosity,
            "cleaned": cleaned,
            "timers": timers,
            "tasks": full_details,
        }

        log_data = {
            "cleaned": cleaned,
            "timers": timers,
        }
        if log_details:
            log_data["tasks"] = log_details

        return full_data, log_data

    def get_report(self, log_message: str, data: dict) -> str:
        """Generate a human-readable report."""

        verbosity: int = data["verbosity"]
        cleaned: bool = data["cleaned"]
        tasks: dict[str, dict[str, Any]] = data["tasks"]

        # Collect task details
        details: list[str] = []
        sum_all_rows: list[tuple[str, ...]] = []
        totals: list[float] = [0, 0, 0.0, 0.0]
        for slug, task in tasks.items():
            # Gather data and types
            summary: dict[str, int] = task["counts"]["summary"]
            needs_cleaning = summary["needs_cleaning"]
            num_cleaned = summary.get("cleaned", 0)

            task_timers: dict[str, float] = task["timers"]
            query_timer = task_timers["query_s"]
            clean_timer = task_timers.get("clean_s", 0.0)

            # Collect summary data
            sum_all_rows.append(
                (
                    slug,
                    str(needs_cleaning),
                    str(num_cleaned),
                    format(query_timer, "0.3f"),
                    format(clean_timer, "0.3f"),
                )
            )
            totals[0] += needs_cleaning
            totals[1] += num_cleaned
            totals[2] += query_timer
            totals[3] += clean_timer

            # Details are ommited on low verbosity
            if verbosity <= 1:
                continue

            # Collect detail section data
            detail = [f"## {slug}"]
            detail.extend(textwrap.wrap(task["check_description"], width=80))
            detail.append("")
            detail.append(
                f"Detected {needs_cleaning} issue{'' if needs_cleaning == 1 else 's'}"
                f" in {query_timer} seconds."
            )
            if cleaned:
                if task["can_clean"]:
                    detail.append(
                        f"Cleaned {num_cleaned} issue{'' if num_cleaned == 1 else 's'}"
                        f" in {clean_timer} seconds."
                    )
                elif needs_cleaning:
                    detail.append("Unable to automatically clean detected items.")
            detail.append("")
            detail.append(task["markdown_report"])
            details.append("\n".join(detail))

        report = ["", "# Summary"]
        if not cleaned:
            report.append("Detected issues only. Use --clean to fix issues.")
        report.append("")

        # Pick summary table columns
        if cleaned:
            columns = ["task", "issues", "fixed", "query (s)", "fix (s)"]
            sum_rows: list[tuple[str, ...]] = sum_all_rows[:]
            sum_rows.append(
                (
                    "_Total_",
                    str(totals[0]),
                    str(totals[1]),
                    format(totals[2], "0.3f"),
                    format(totals[3], "0.3f"),
                )
            )
        else:
            columns = ["task", "issues", "query (s)"]
            sum_rows = [(row[0], row[1], row[3]) for row in sum_all_rows]
            sum_rows.append(("_Total_", str(totals[0]), format(totals[2], "0.3f")))

        # Determine summary table widths
        widths = [len(col) for col in columns]
        for row in sum_rows:
            for column, value in enumerate(row):
                widths[column] = max(widths[column], len(value))

        # Output summary table with aligned columns
        header = (
            "|"
            + "|".join(f"{col:^{widths[colnum]}}" for colnum, col in enumerate(columns))
            + "|"
        )
        sep = (
            "|"
            + "|".join(
                f"{'-' * (widths[colnum] - 1)}:" for colnum, _ in enumerate(columns)
            )
            + "|"
        )
        report.extend([header, sep])
        for row in sum_rows:
            row_report = (
                "|"
                + "|".join(f"{col:>{widths[colnum]}}" for colnum, col in enumerate(row))
                + "|"
            )
            report.append(row_report)

        # Output details
        if verbosity > 1:
            report.extend(["", "# Details"])
            for detail_block in details:
                report.extend([detail_block, ""])

        return "\n".join(report)

1	from __future__ import annotations	1✔
2
3	import logging	1✔
4	import textwrap	1✔
5	from argparse import RawDescriptionHelpFormatter	1✔
6	from shutil import get_terminal_size	1✔
7	from typing import TYPE_CHECKING, Any	1✔
8
9	from django.core.management.base import BaseCommand, DjangoHelpFormatter	1✔
10	from django.db import transaction	1✔
11
12	from codetiming import Timer	1✔
13
14	from emails.cleaners import MissingProfileCleaner, ServerStorageCleaner	1✔
15	from privaterelay.cleaners import MissingEmailCleaner	1✔
16
17	if TYPE_CHECKING: # pragma: no cover
18	from argparse import ArgumentParser
19
20	from privaterelay.cleaner_task import DataIssueTask
21
22
23	logger = logging.getLogger("eventsinfo.cleanup_data")	1✔
24
25
26	class RawDescriptionDjangoHelpFormatter(	1✔
27	DjangoHelpFormatter, RawDescriptionHelpFormatter
28	):
29	"""DjangoHelpFormatter, but don't reflow the epilog."""
30
31
32	class Command(BaseCommand):	1✔
33	help = "Detects and optionally clean data issues."	1✔
34	task_list: list[type[DataIssueTask]] = [	1✔
35	ServerStorageCleaner,
36	MissingProfileCleaner,
37	MissingEmailCleaner,
38	]
39	tasks: dict[str, DataIssueTask]	1✔
40
41	def __init__(self, args, *kwargs):	1✔
42	super().__init__(args, *kwargs)	1✔
43	tasks = {Task.slug: Task() for Task in self.task_list}	1✔
44	self.tasks = {slug: tasks[slug] for slug in sorted(tasks.keys())}	1✔
45
46	def create_parser(self, prog_name, subcommand, **kwargs):	1✔
47	"""Add DataIssueTask information to default parser."""
48
49	epilog_lines = ["Data issue detectors / cleaners:"]	1✔
50	for slug, task in self.tasks.items():	1✔
51	raw = f"{slug} - {task.title}"	1✔
52	epilog_lines.extend(	1✔
53	textwrap.wrap(
54	raw,
55	width=get_terminal_size().columns,
56	initial_indent=" ",
57	subsequent_indent=" ",
58	)
59	)
60	epilog = "\n".join(epilog_lines)	1✔
61
62	parser = super().create_parser(prog_name, subcommand, epilog=epilog, **kwargs)	1✔
63	if not parser.formatter_class == DjangoHelpFormatter:	1!
64	raise TypeError("parser.formatter_class must be type DjangoHelpFormatter")	×
65	parser.formatter_class = RawDescriptionDjangoHelpFormatter	1✔
66	return parser	1✔
67
68	def add_arguments(self, parser: ArgumentParser) -> None:	1✔
69	parser.add_argument(	1✔
70	"--clean", action="store_true", help="Clean detected data issues."
71	)
72	for slug, task in self.tasks.items():	1✔
73	parser.add_argument(	1✔
74	f"--{slug}",
75	action="store_true",
76	help=f"Only run {slug} {'cleaner' if task.can_clean else 'detector'}",
77	)
78
79	def handle(self, clean: bool, verbosity: int, args: Any, *kwargs: Any) -> str:	1✔
80	"""Run the cleanup_data command."""
81
82	# Determine if we're running some tasks or the full set
83	run_some: list[str] = []	1✔
84	for name, val in kwargs.items():	1✔
85	slug = name.replace("_", "-")	1✔
86	if slug in self.tasks and val:	1✔
87	run_some.append(slug)	1✔
88	if run_some:	1✔
89	tasks = {slug: self.tasks[slug] for slug in sorted(run_some)}	1✔
90	else:
91	tasks = self.tasks	1✔
92
93	# Find data issues and clean them if requested
94	issue_count = 0	1✔
95	clean_count = 0	1✔
96	issue_timers: dict[str, float] = {}	1✔
97	clean_timers: dict[str, float] = {}	1✔
98	for slug, task in tasks.items():	1✔
99	with transaction.atomic():	1✔
100	task_issues, task_time = self.find_issues_in_task(task)	1✔
101	issue_count += task_issues	1✔
102	issue_timers[slug] = task_time	1✔
103
104	if clean:	1✔
105	task_cleaned, clean_time = self.clean_issues_in_task(task)	1✔
106	if task_cleaned is not None and clean_time is not None:	1!
107	clean_count += task_cleaned	1✔
108	clean_timers[slug] = clean_time	1✔
109
110	# Log results and create a report, based on requested verbosity
111	if clean:	1✔
112	log_message = (	1✔
113	f"cleanup_data complete, cleaned {clean_count} of"
114	f" {issue_count} issue{'' if issue_count==1 else 's'}."
115	)
116	else:
117	log_message = (	1✔
118	f"cleanup_data complete, found {issue_count}"
119	f" issue{'' if issue_count==1 else 's'} (dry run)."
120	)
121
122	full_data, log_data = self.prepare_data(	1✔
123	clean, verbosity, tasks, issue_timers, clean_timers=clean_timers
124	)
125	logger.info(log_message, extra=log_data)	1✔
126	report = self.get_report(log_message, full_data)	1✔
127	return report	1✔
128
129	def find_issues_in_task(self, task: DataIssueTask) -> tuple[int, float]:	1✔
130	"""Run a task, returning the number of issues and execution time."""
131	with Timer(logger=None) as issue_timer:	1✔
132	total = task.issues()	1✔
133	return total, round(issue_timer.last, 3)	1✔
134
135	def clean_issues_in_task(	1✔
136	self, task: DataIssueTask
137	) -> tuple[int, float] \| tuple[None, None]:
138	"""Run the task cleaner, returning issues cleaned and execution time."""
139	if not hasattr(task, "clean"):	1!
NEW 140	return (None, None)	×
141	with Timer(logger=None) as clean_timer:	1✔
142	cleaned = task.clean()	1✔
143	return cleaned, round(clean_timer.last, 3)	1✔
144
145	def prepare_data(	1✔
146	self,
147	cleaned: bool,
148	verbosity: int,
149	tasks: dict[str, DataIssueTask],
150	issue_timers: dict[str, float],
151	clean_timers: dict[str, float],
152	) -> tuple[dict, dict]:
153	"""Gather full data and log data from tasks and timers."""
154
155	timers = {"query_s": sum(issue_timers.values())}	1✔
156	if clean_timers:	1✔
157	timers["clean_s"] = sum(clean_timers.values())	1✔
158
159	log_details = {}	1✔
160	full_details = {}	1✔
161	for slug, task in tasks.items():	1✔
162	timers = {"query_s": issue_timers[slug]}	1✔
163	if clean_timers:	1✔
164	timers["clean_s"] = clean_timers[slug]	1✔
165
166	full_details[slug] = {	1✔
167	"title": task.title,
168	"check_description": task.check_description,
169	"can_clean": task.can_clean,
170	"counts": task.counts,
171	"markdown_report": task.markdown_report(),
172	"timers": timers,
173	}
174
175	if verbosity > 1:	1✔
176	log_details[slug] = {"counts": task.counts, "timers": timers}	1✔
177	elif verbosity > 0:	1✔
178	log_details[slug] = {"counts": {"summary": task.counts["summary"]}}	1✔
179
180	full_data = {	1✔
181	"verbosity": verbosity,
182	"cleaned": cleaned,
183	"timers": timers,
184	"tasks": full_details,
185	}
186
187	log_data = {	1✔
188	"cleaned": cleaned,
189	"timers": timers,
190	}
191	if log_details:	1✔
192	log_data["tasks"] = log_details	1✔
193
194	return full_data, log_data	1✔
195
196	def get_report(self, log_message: str, data: dict) -> str:	1✔
197	"""Generate a human-readable report."""
198
199	verbosity: int = data["verbosity"]	1✔
200	cleaned: bool = data["cleaned"]	1✔
201	tasks: dict[str, dict[str, Any]] = data["tasks"]	1✔
202
203	# Collect task details
204	details: list[str] = []	1✔
205	sum_all_rows: list[tuple[str, ...]] = []	1✔
206	totals: list[float] = [0, 0, 0.0, 0.0]	1✔
207	for slug, task in tasks.items():	1✔
208	# Gather data and types
209	summary: dict[str, int] = task["counts"]["summary"]	1✔
210	needs_cleaning = summary["needs_cleaning"]	1✔
211	num_cleaned = summary.get("cleaned", 0)	1✔
212
213	task_timers: dict[str, float] = task["timers"]	1✔
214	query_timer = task_timers["query_s"]	1✔
215	clean_timer = task_timers.get("clean_s", 0.0)	1✔
216
217	# Collect summary data
218	sum_all_rows.append(	1✔
219	(
220	slug,
221	str(needs_cleaning),
222	str(num_cleaned),
223	format(query_timer, "0.3f"),
224	format(clean_timer, "0.3f"),
225	)
226	)
227	totals[0] += needs_cleaning	1✔
228	totals[1] += num_cleaned	1✔
229	totals[2] += query_timer	1✔
230	totals[3] += clean_timer	1✔
231
232	# Details are ommited on low verbosity
233	if verbosity <= 1:	1✔
234	continue	1✔
235
236	# Collect detail section data
237	detail = [f"## {slug}"]	1✔
238	detail.extend(textwrap.wrap(task["check_description"], width=80))	1✔
239	detail.append("")	1✔
240	detail.append(	1✔
241	f"Detected {needs_cleaning} issue{'' if needs_cleaning == 1 else 's'}"
242	f" in {query_timer} seconds."
243	)
244	if cleaned:	1✔
245	if task["can_clean"]:	1!
246	detail.append(	1✔
247	f"Cleaned {num_cleaned} issue{'' if num_cleaned == 1 else 's'}"
248	f" in {clean_timer} seconds."
249	)
250	elif needs_cleaning:	×
251	detail.append("Unable to automatically clean detected items.")	×
252	detail.append("")	1✔
253	detail.append(task["markdown_report"])	1✔
254	details.append("\n".join(detail))	1✔
255
256	report = ["", "# Summary"]	1✔
257	if not cleaned:	1✔
258	report.append("Detected issues only. Use --clean to fix issues.")	1✔
259	report.append("")	1✔
260
261	# Pick summary table columns
262	if cleaned:	1✔
263	columns = ["task", "issues", "fixed", "query (s)", "fix (s)"]	1✔
264	sum_rows: list[tuple[str, ...]] = sum_all_rows[:]	1✔
265	sum_rows.append(	1✔
266	(
267	"_Total_",
268	str(totals[0]),
269	str(totals[1]),
270	format(totals[2], "0.3f"),
271	format(totals[3], "0.3f"),
272	)
273	)
274	else:
275	columns = ["task", "issues", "query (s)"]	1✔
276	sum_rows = [(row[0], row[1], row[3]) for row in sum_all_rows]	1✔
277	sum_rows.append(("_Total_", str(totals[0]), format(totals[2], "0.3f")))	1✔
278
279	# Determine summary table widths
280	widths = [len(col) for col in columns]	1✔
281	for row in sum_rows:	1✔
282	for column, value in enumerate(row):	1✔
283	widths[column] = max(widths[column], len(value))	1✔
284
285	# Output summary table with aligned columns
286	header = (	1✔
287	"\|"
288	+ "\|".join(f"{col:^{widths[colnum]}}" for colnum, col in enumerate(columns))
289	+ "\|"
290	)
291	sep = (	1✔
292	"\|"
293	+ "\|".join(
294	f"{'-' * (widths[colnum] - 1)}:" for colnum, _ in enumerate(columns)
295	)
296	+ "\|"
297	)
298	report.extend([header, sep])	1✔
299	for row in sum_rows:	1✔
300	row_report = (	1✔
301	"\|"
302	+ "\|".join(f"{col:>{widths[colnum]}}" for colnum, col in enumerate(row))
303	+ "\|"
304	)
305	report.append(row_report)	1✔
306
307	# Output details
308	if verbosity > 1:	1✔
309	report.extend(["", "# Details"])	1✔
310	for detail_block in details:	1✔
311	report.extend([detail_block, ""])	1✔
312
313	return "\n".join(report)	1✔

mozilla / fx-private-relay / 53a4fc85-bb39-4dee-81e5-c2abd320ed53

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous