python-threatexchange/threatexchange/cli/fetch_cmd.py

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import collections
import logging
import datetime
import logging
import time
import typing as t

from threatexchange.cli.cli_config import CLISettings
from threatexchange.cli.dataset_cmd import DatasetCommand
from threatexchange.fetcher.collab_config import CollaborationConfigBase
from threatexchange.fetcher.fetch_api import SignalExchangeAPI
from threatexchange.fetcher.fetch_state import (
    FetchCheckpointBase,
    FetchedStateStoreBase,
)
from threatexchange.cli import command_base


class FetchCommand(command_base.Command):
    """
    Download content from signal exchange APIs to disk.
    """

    PROGRESS_PRINT_INTERVAL_SEC = 30

    @classmethod
    def init_argparse(cls, settings: CLISettings, ap) -> None:
        ap.add_argument(
            "--clear",
            action="store_true",
            help="delete fetched state and checkpoints "
            "(you almost never need to do this)",
        )
        ap.add_argument(
            "--skip-index-rebuild",
            action="store_true",
            help="don't rebuild indices after fetch",
        )
        ap.add_argument("--limit", type=int, help="stop after fetching this many items")
        ap.add_argument(
            "--time-limit-sec",
            type=int,
            metavar="SEC",
            help="stop fetching after this many seconds",
        )
        ap.add_argument(
            "--only-api",
            choices=[f.get_name() for f in settings.get_fetchers()],
            help="only fetch from this API",
        )
        ap.add_argument(
            "--only-collab",
            metavar="NAME",
            help="only fetch for this collaboration",
        )

    def __init__(
        self,
        # Defaults to make it easier to call from match
        clear: bool = False,
        time_limit_sec: t.Optional[int] = None,
        limit: t.Optional[int] = None,
        skip_index_rebuild: bool = False,
        only_api: t.Optional[str] = None,
        only_collab: t.Optional[str] = None,
    ) -> None:
        self.clear = clear
        self.time_limit_sec = time_limit_sec
        self.limit = limit
        self.skip_index_rebuild = skip_index_rebuild
        self.only_api = only_api
        self.only_collab = only_collab
        self.collabs: t.List[CollaborationConfigBase] = []

        # Limits
        self.total_fetched_count = 0
        self.start_time = time.time()

        # Progress
        self.last_update_time: t.Optional[int] = None
        # Print first update after 5 seconds
        self.last_update_printed = time.time() - self.PROGRESS_PRINT_INTERVAL_SEC + 5
        self.progress_fetched_count = 0
        self.counts: t.Dict[str, int] = collections.Counter()

    def has_hit_limits(self):
        if self.limit is not None and self.total_fetched_count >= self.limit:
            return True
        if self.time_limit_sec is not None:
            if time.time() - self.start_time >= self.time_limit_sec:
                return True
        return False

    def execute(self, settings: CLISettings) -> None:
        fetchers = settings.get_fetchers()
        # Verify collab arguments
        self.collabs = settings.get_all_collabs(default_to_sample=True)
        if self.only_collab:
            self.collabs = [c for c in self.collabs if c.name == self.only_collab]
            if not self.collabs:
                raise command_base.CommandError(
                    f"No such collab '{self.only_collab}'", 2
                )
        if all(not c.enabled for c in self.collabs):
            self.stderr("All collabs are disabled. Nothing to do.")
            return

        # Do work
        if self.clear:
            self.stderr("Clearing fetched state")
            for fetcher in settings.get_fetchers():
                store = settings.get_fetch_store_for_fetcher(fetcher)
                for collab in self.collabs:
                    if self.only_collab not in (None, collab.name):
                        continue
                    logging.info("Clearing %s - %s", fetcher.get_name(), collab.name)
                    store.clear(collab)
            return

        all_succeeded = True
        any_succeded = False

        for fetcher in fetchers:
            logging.info("Fetching all %s's configs", fetcher.get_name())
            succeeded = self.execute_for_fetcher(settings, fetcher)
            all_succeeded &= succeeded
            any_succeded |= succeeded

        if any_succeded and not self.skip_index_rebuild:
            self.stderr("Rebuilding match indices...")
            DatasetCommand().execute_generate_indices(settings)

        if not all_succeeded:
            raise command_base.CommandError("Some collabs had errors!", 3)

    def execute_for_fetcher(
        self, settings: CLISettings, fetcher: SignalExchangeAPI
    ) -> bool:
        success = True
        for collab in self.collabs:
            if collab.api != fetcher.get_name():
                continue
            if not collab.enabled:
                logging.debug(
                    "Skipping %s, disabled",
                )
                continue
            fetch_ok = self.execute_for_collab(settings, fetcher, collab)
            success &= fetch_ok
        return success

    def execute_for_collab(
        self,
        settings: CLISettings,
        fetcher: SignalExchangeAPI,
        collab: CollaborationConfigBase,
    ) -> bool:

        store = settings.get_fetch_store_for_fetcher(fetcher.__class__)
        checkpoint = self._verify_store_and_checkpoint(store, collab)

        self.progress_fetched_count = 0
        self.current_collab = collab.name
        self.current_api = fetcher.get_name()

        try:
            while not self.has_hit_limits():
                delta = fetcher.fetch_once(
                    settings.get_all_signal_types(), collab, checkpoint
                )
                logging.info("Fetched %d records", delta.record_count())
                checkpoint = delta.next_checkpoint()
                self._fetch_progress(delta.record_count(), checkpoint)
                assert checkpoint is not None  # Infinite loop protection
                store.merge(collab, delta)
                if not delta.has_more():
                    break
        except:
            self._stderr_current("failed to fetch!")
            logging.exception("Failed to fetch %s", collab.name)
            return False
        finally:
            store.flush()

        self._print_progress(done=True)
        return True

    def _verify_store_and_checkpoint(
        self, store: FetchedStateStoreBase, collab: CollaborationConfigBase
    ) -> t.Optional[FetchCheckpointBase]:
        checkpoint = store.get_checkpoint(collab)

        if checkpoint is not None and checkpoint.is_stale():
            store.clear(collab)
            return None

        return checkpoint

    def _fetch_progress(self, batch_size: int, checkpoint: FetchCheckpointBase) -> None:
        self.progress_fetched_count += batch_size
        self.total_fetched_count += batch_size
        progress_ts = checkpoint.get_progress_timestamp()
        if progress_ts is not None:
            self.last_update_time = progress_ts

        now = time.time()
        if now - self.last_update_printed >= self.PROGRESS_PRINT_INTERVAL_SEC:
            self.last_update_printed = now
            self._print_progress()

    def _stderr_current(self, msg: str) -> None:
        assert self.current_api and self.current_collab
        self.stderr(
            f"[{self.current_api}] {self.current_collab} - {msg}",
        )

    def _print_progress(self, *, done=False):
        processed = "Syncing..."
        if done:
            processed = "Up to date"
        elif self.progress_fetched_count:
            processed = f"Downloaded {self.progress_fetched_count} updates"

        from_time = ""
        if self.last_update_time is not None:
            if not from_time:
                from_time = "ages long past"
            elif self.last_update_time >= time.time() - 1:
                from_time = "moments ago"
            else:
                from_time = datetime.datetime.fromtimestamp(
                    self.last_update_time
                ).isoformat()
            from_time = f", at {from_time}"

        self._stderr_current(f"{processed}{from_time}")