-
Notifications
You must be signed in to change notification settings - Fork 304
/
helpers.py
147 lines (118 loc) · 4.45 KB
/
helpers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# Copyright (c) Meta Platforms, Inc. and affiliates.
"""
Helpers for implementing SignalExchangeAPIs or state.
They demonstrate how the methods in SignalExchangeAPI can be used to
reproduce databases of signals, as well as simple recipes for stage.
"""
from dataclasses import dataclass, field
import logging
import typing as t
from threatexchange.exchanges.signal_exchange_api import TSignalExchangeAPICls
from threatexchange.signal_type.signal_base import SignalType
from threatexchange.exchanges import fetch_state
from threatexchange.exchanges.collab_config import CollaborationConfigBase
K = t.TypeVar("K")
V = t.TypeVar("V")
@dataclass
class SimpleFetchedSignalMetadata(fetch_state.FetchedSignalMetadata):
"""
Simple dataclass for fetched data.
Merge by addition rather than replacement.
"""
opinions: t.List[fetch_state.SignalOpinion] = field(default_factory=list)
def get_as_opinions(self) -> t.Sequence[fetch_state.SignalOpinion]:
return self.opinions
@dataclass
class _StateTracker:
api_cls: TSignalExchangeAPICls
_delta: t.Optional[fetch_state.FetchDeltaTyped]
dirty: bool = False
@property
def empty(self) -> bool:
return self._delta is None
@property
def delta(self) -> fetch_state.FetchDeltaTyped:
assert self._delta is not None
return self._delta
@delta.setter
def delta(self, value: fetch_state.FetchDeltaTyped) -> None:
if self._delta is None:
old = {}
self._delta = value
else:
old = self._delta.updates
self.api_cls.naive_fetch_merge(old, value.updates)
self._delta.updates = old
self._delta.checkpoint = value.checkpoint
self.dirty = True
@property
def checkpoint(self) -> t.Optional[fetch_state.FetchCheckpointBase]:
return None if self._delta is None else self._delta.checkpoint
class SimpleFetchedStateStore(fetch_state.FetchedStateStoreBase):
"""
A FetchedStateStore that does merges in memory and writes by collab
"""
def __init__(
self,
api_cls: TSignalExchangeAPICls,
) -> None:
self.api_cls = api_cls
self._state: t.Dict[str, _StateTracker] = {}
def _read_state(
self,
collab_name: str,
) -> t.Optional[fetch_state.FetchDeltaTyped]:
raise NotImplementedError
def _write_state(
self, collab_name: str, delta: fetch_state.FetchDeltaTyped
) -> None:
raise NotImplementedError
def get_checkpoint(
self, collab: CollaborationConfigBase
) -> t.Optional[fetch_state.FetchCheckpointBase]:
return self._get_state(collab).checkpoint
def _get_state(self, collab: CollaborationConfigBase) -> _StateTracker:
if collab.name not in self._state:
logging.debug("Loading state for %s", collab.name)
delta = self._read_state(collab.name)
self._state[collab.name] = _StateTracker(self.api_cls, delta)
return self._state[collab.name]
def clear(self, collab: CollaborationConfigBase) -> None:
self._state.pop(collab.name, None)
def merge(
self,
collab: CollaborationConfigBase,
delta: fetch_state.FetchDeltaTyped,
) -> None:
"""
Merge a FetchDeltaBase into the state.
At the implementation's discretion, it may call flush() or the
equivalent work.
"""
state = self._get_state(collab)
if len(delta.updates) == 0 and delta.checkpoint in (
None,
state.checkpoint,
):
logging.warning("No op update for %s", collab.name)
return
state.delta = delta
def flush(self) -> None:
for collab_name, state in self._state.items():
if state.dirty:
assert state.delta is not None
self._write_state(collab_name, state.delta)
state.dirty = False
def get_for_signal_type(
self, collabs: t.List[CollaborationConfigBase], signal_type: t.Type[SignalType]
) -> t.Dict[str, t.Dict[str, fetch_state.FetchedSignalMetadata]]:
ret = {}
for collab in collabs:
state = self._get_state(collab)
if not state.empty:
by_signal = state.api_cls.naive_convert_to_signal_type(
[signal_type], collab, state.delta.updates
).get(signal_type, {})
if by_signal:
ret[collab.name] = by_signal
return ret