Changeset View
Changeset View
Standalone View
Standalone View
swh/indexer/storage/in_memory.py
| Show All 35 Lines | from .model import ( | ||||
| ContentCtagsRow, | ContentCtagsRow, | ||||
| ContentLanguageRow, | ContentLanguageRow, | ||||
| ContentLicenseRow, | ContentLicenseRow, | ||||
| ContentMetadataRow, | ContentMetadataRow, | ||||
| ContentMimetypeRow, | ContentMimetypeRow, | ||||
| OriginIntrinsicMetadataRow, | OriginIntrinsicMetadataRow, | ||||
| RevisionIntrinsicMetadataRow, | RevisionIntrinsicMetadataRow, | ||||
| ) | ) | ||||
| from .writer import JournalWriter | |||||
| SHA1_DIGEST_SIZE = 160 | SHA1_DIGEST_SIZE = 160 | ||||
| ToolId = int | |||||
| def _transform_tool(tool): | def _transform_tool(tool): | ||||
| return { | return { | ||||
| "id": tool["id"], | "id": tool["id"], | ||||
| "name": tool["tool_name"], | "name": tool["tool_name"], | ||||
| "version": tool["tool_version"], | "version": tool["tool_version"], | ||||
| "configuration": tool["tool_configuration"], | "configuration": tool["tool_configuration"], | ||||
| } | } | ||||
| def check_id_types(data: List[Dict[str, Any]]): | def check_id_types(data: List[Dict[str, Any]]): | ||||
| """Checks all elements of the list have an 'id' whose type is 'bytes'.""" | """Checks all elements of the list have an 'id' whose type is 'bytes'.""" | ||||
| if not all(isinstance(item.get("id"), bytes) for item in data): | if not all(isinstance(item.get("id"), bytes) for item in data): | ||||
| raise IndexerStorageArgumentException("identifiers must be bytes.") | raise IndexerStorageArgumentException("identifiers must be bytes.") | ||||
| def _key_from_dict(d): | def _key_from_dict(d): | ||||
| return tuple(sorted(d.items())) | return tuple(sorted(d.items())) | ||||
| ToolId = int | |||||
| TValue = TypeVar("TValue", bound=BaseRow) | TValue = TypeVar("TValue", bound=BaseRow) | ||||
| class SubStorage(Generic[TValue]): | class SubStorage(Generic[TValue]): | ||||
| """Implements common missing/get/add logic for each indexer type.""" | """Implements common missing/get/add logic for each indexer type.""" | ||||
| _data: Dict[Sha1, Dict[Tuple, Dict[str, Any]]] | _data: Dict[Sha1, Dict[Tuple, Dict[str, Any]]] | ||||
| _tools_per_id: Dict[Sha1, Set[ToolId]] | _tools_per_id: Dict[Sha1, Set[ToolId]] | ||||
| def __init__(self, row_class: Type[TValue], tools): | def __init__(self, row_class: Type[TValue], tools, journal_writer): | ||||
| self.row_class = row_class | self.row_class = row_class | ||||
| self._tools = tools | self._tools = tools | ||||
| self._sorted_ids = SortedList[bytes, Sha1]() | self._sorted_ids = SortedList[bytes, Sha1]() | ||||
| self._data = defaultdict(dict) | self._data = defaultdict(dict) | ||||
| self._journal_writer = journal_writer | |||||
| self._tools_per_id = defaultdict(set) | self._tools_per_id = defaultdict(set) | ||||
| def _key_from_dict(self, d) -> Tuple: | def _key_from_dict(self, d) -> Tuple: | ||||
| """Like the global _key_from_dict, but filters out dict keys that don't | """Like the global _key_from_dict, but filters out dict keys that don't | ||||
| belong in the unique key.""" | belong in the unique key.""" | ||||
| return _key_from_dict({k: d[k] for k in self.row_class.UNIQUE_KEY_FIELDS}) | return _key_from_dict({k: d[k] for k in self.row_class.UNIQUE_KEY_FIELDS}) | ||||
| def missing(self, keys: Iterable[Dict]) -> List[Sha1]: | def missing(self, keys: Iterable[Dict]) -> List[Sha1]: | ||||
| ▲ Show 20 Lines • Show All 112 Lines • ▼ Show 20 Lines | def add(self, data: Iterable[TValue]) -> int: | ||||
| - **id**: sha1 | - **id**: sha1 | ||||
| - **indexer_configuration_id**: tool used to compute the | - **indexer_configuration_id**: tool used to compute the | ||||
| results | results | ||||
| - arbitrary data | - arbitrary data | ||||
| """ | """ | ||||
| data = list(data) | data = list(data) | ||||
| check_id_duplicates(data) | check_id_duplicates(data) | ||||
| object_type = self.row_class.object_type # type: ignore | |||||
| self._journal_writer.write_additions(object_type, data) | |||||
| count = 0 | count = 0 | ||||
| for obj in data: | for obj in data: | ||||
| item = obj.to_dict() | item = obj.to_dict() | ||||
| id_ = item.pop("id") | id_ = item.pop("id") | ||||
| tool_id = item["indexer_configuration_id"] | tool_id = item["indexer_configuration_id"] | ||||
| key = _key_from_dict(obj.unique_key()) | key = _key_from_dict(obj.unique_key()) | ||||
| self._data[id_][key] = item | self._data[id_][key] = item | ||||
| self._tools_per_id[id_].add(tool_id) | self._tools_per_id[id_].add(tool_id) | ||||
| count += 1 | count += 1 | ||||
| if id_ not in self._sorted_ids: | if id_ not in self._sorted_ids: | ||||
| self._sorted_ids.add(id_) | self._sorted_ids.add(id_) | ||||
| return count | return count | ||||
| class IndexerStorage: | class IndexerStorage: | ||||
| """In-memory SWH indexer storage.""" | """In-memory SWH indexer storage.""" | ||||
| def __init__(self): | def __init__(self, journal_writer=None): | ||||
| self._tools = {} | self._tools = {} | ||||
| self._mimetypes = SubStorage(ContentMimetypeRow, self._tools) | |||||
| self._languages = SubStorage(ContentLanguageRow, self._tools) | def tool_getter(id_): | ||||
| self._content_ctags = SubStorage(ContentCtagsRow, self._tools) | tool = self._tools[id_] | ||||
| self._licenses = SubStorage(ContentLicenseRow, self._tools) | return { | ||||
| self._content_metadata = SubStorage(ContentMetadataRow, self._tools) | "id": tool["id"], | ||||
| "name": tool["tool_name"], | |||||
| "version": tool["tool_version"], | |||||
| "configuration": tool["tool_configuration"], | |||||
| } | |||||
| self.journal_writer = JournalWriter(tool_getter, journal_writer) | |||||
| args = (self._tools, self.journal_writer) | |||||
| self._mimetypes = SubStorage(ContentMimetypeRow, *args) | |||||
| self._languages = SubStorage(ContentLanguageRow, *args) | |||||
| self._content_ctags = SubStorage(ContentCtagsRow, *args) | |||||
| self._licenses = SubStorage(ContentLicenseRow, *args) | |||||
| self._content_metadata = SubStorage(ContentMetadataRow, *args) | |||||
| self._revision_intrinsic_metadata = SubStorage( | self._revision_intrinsic_metadata = SubStorage( | ||||
| RevisionIntrinsicMetadataRow, self._tools | RevisionIntrinsicMetadataRow, *args | ||||
| ) | |||||
| self._origin_intrinsic_metadata = SubStorage( | |||||
| OriginIntrinsicMetadataRow, self._tools | |||||
| ) | ) | ||||
| self._origin_intrinsic_metadata = SubStorage(OriginIntrinsicMetadataRow, *args) | |||||
| def check_config(self, *, check_write): | def check_config(self, *, check_write): | ||||
| return True | return True | ||||
| def content_mimetype_missing( | def content_mimetype_missing( | ||||
| self, mimetypes: Iterable[Dict] | self, mimetypes: Iterable[Dict] | ||||
| ) -> List[Tuple[Sha1, int]]: | ) -> List[Tuple[Sha1, int]]: | ||||
| return self._mimetypes.missing(mimetypes) | return self._mimetypes.missing(mimetypes) | ||||
| ▲ Show 20 Lines • Show All 239 Lines • Show Last 20 Lines | |||||