def test_grouper_with_stop_value(self):
# given
actual_data = utils.grouper(((i, i + 1) for i in range(0, 9)), 2)
out = []
for d in actual_data:
out.append(list(d)) # force generator resolution for checks
self.assertEqual(
out,
[
[(0, 1), (1, 2)],
[(2, 3), (3, 4)],
[(4, 5), (5, 6)],
[(6, 7), (7, 8)],
[(8, 9)],
],
)
# given
actual_data = utils.grouper((i for i in range(9, 0, -1)), 4)
out = []
for d in actual_data:
out.append(list(d)) # force generator resolution for checks
self.assertEqual(out, [[9, 8, 7, 6], [5, 4, 3, 2], [1]])
def random_blocks(iterable, block=100):
"""Randomize iterable per block of size block.
Given an iterable:
- slice the iterable in data set of block-sized elements
- randomized the block-sized elements
- yield each element of that randomized block-sized
- continue onto the next block-sized block
Args:
iterable (Iterable): an iterable
block (int): number of elements per block
Yields:
random element of the iterable
"""
count = 0
for iter_ in grouper(iterable, block):
count += 1
lst = list(iter_)
random.shuffle(lst)
for e in lst:
yield e
def add(self, swhids: Iterable[str], chunk_size: int, cur: sqlite3.Cursor):
"""Insert the SWHID inside the database."""
for swhids_chunk in grouper(swhids, chunk_size):
cur.executemany(
"""INSERT INTO swhids VALUES (?)""",
[(swhid_chunk, ) for swhid_chunk in swhids_chunk],
)
def ingest_data(self, identifier, checks=False):
"""Rework the base ingest_data.
Request server endpoint which gives all in one go.
Simplify and filter response list of repositories. Inject
repo information into local db. Queue loader tasks for
linked repositories.
Args:
identifier: Resource identifier (unused)
checks (bool): Additional checks required (unused)
"""
response = self.safely_issue_request(identifier)
response = self.list_packages(response)
if not response:
return response, []
models_list = self.transport_response_simplified(response)
models_list = self.filter_before_inject(models_list)
all_injected = []
for models in utils.grouper(models_list, n=10000):
models = list(models)
logging.debug('models: %s' % len(models))
# inject into local db
injected = self.inject_repo_data_into_db(models)
# queue workers
self.create_missing_origins_and_tasks(models, injected)
all_injected.append(injected)
# flush
self.db_session.commit()
self.db_session = self.mk_session()
return response, all_injected
def call_with_batches(
f: Callable[[List[T1]], Iterable[T2]], args: List[T1], batch_size: int,
) -> Iterator[T2]:
"""Calls a function with batches of args, and concatenates the results.
"""
groups = grouper(args, batch_size)
for group in groups:
yield from f(list(group))
def retrieve_unknown_sha1s(swhstorage, gen_data):
# Compute blocks of 1000 sha1s
for paths in grouper(gen_data, n=1000):
data = process_paths(paths)
sha1s_tocheck = list(data.keys())
if len(sha1s_tocheck) > 0:
# let those inexistent sha1s flow
for sha1 in swhstorage.content_missing_per_sha1(sha1s_tocheck):
yield data[sha1][0], data[sha1][1]
def test_grouper():
# given
actual_data = utils.grouper((i for i in range(0, 9)), 2)
out = []
for d in actual_data:
out.append(list(d)) # force generator resolution for checks
assert out == [[0, 1], [2, 3], [4, 5], [6, 7], [8]]
# given
actual_data = utils.grouper((i for i in range(9, 0, -1)), 4)
out = []
for d in actual_data:
out.append(list(d)) # force generator resolution for checks
assert out == [[9, 8, 7, 6], [5, 4, 3, 2], [1]]
def retrieve_unknown_sha1s(swhstorage, gen_data):
# Compute blocks of 1000 sha1s
for paths in grouper(gen_data, n=1000):
data = process_paths(paths)
sha1s_tocheck = list(data.keys())
if len(sha1s_tocheck) > 0:
# let those inexistent sha1s flow
for sha1 in swhstorage.content_missing_per_sha1(sha1s_tocheck):
yield data[sha1][0], data[sha1][1]
def store_file_content(db_url, path):
"""The first round finished, there were errors. Adapting the code and
running this command will finish appropriately the first round.
"""
db = Db.connect(db_url)
with db.transaction() as cur:
for data in grouper(load_file(path), BLOCK_SIZE):
db.copy_to(data, 'content_sesi',
['origin_sha1', 'sha1', 'sha1_git', 'sha256', 'length',
'path', 'corrupted'], cur)
def get_new_contents_metadata(
self, all_contents: List[Dict[str, Any]]
) -> Generator[Tuple[Dict[str, Any], List[Any]], Any, None]:
"""Retrieve raw contents and compute new checksums on the
contents. Unknown or corrupted contents are skipped.
Args:
all_contents: List of contents as dictionary with
the necessary primary keys
Yields:
tuple: tuple of (content to update, list of checksums computed)
"""
content_ids = self._read_content_ids(all_contents)
for contents in utils.grouper(content_ids,
self.batch_size_retrieve_content):
contents_iter = itertools.tee(contents, 2)
try:
sha1s = [s for s in contents_iter[0]]
content_metadata: List[
Optional[Content]] = self.storage.content_get(sha1s)
except Exception:
self.log.exception("Problem when reading contents metadata.")
continue
for sha1, content_model in zip(sha1s, content_metadata):
if not content_model:
continue
content: Dict = content_model.to_dict()
# Recompute checksums provided in compute_checksums options
if self.recompute_checksums:
checksums_to_compute = list(self.compute_checksums)
else:
# Compute checksums provided in compute_checksums
# options not already defined for that content
checksums_to_compute = [
h for h in self.compute_checksums if not content.get(h)
]
if not checksums_to_compute: # Nothing to recompute
continue
try:
raw_content = self.objstorage.get(sha1)
except ObjNotFoundError:
self.log.warning("Content %s not found in objstorage!",
sha1)
continue
content_hashes = hashutil.MultiHash.from_data(
raw_content, hash_names=checksums_to_compute).digest()
content.update(content_hashes)
yield content, checksums_to_compute
def store_file_content(db_url, path):
"""The first round finished, there were errors. Adapting the code and
running this command will finish appropriately the first round.
"""
db = Db.connect(db_url)
with db.transaction() as cur:
for data in grouper(load_file(path), BLOCK_SIZE):
db.copy_to(data, 'content_sesi', [
'origin_sha1', 'sha1', 'sha1_git', 'sha256', 'length', 'path',
'corrupted'
], cur)
def copy_finished(self, content_ids):
"""Once the copy is finished, we'll send those batch of contents as
done in the destination queue.
"""
if self.task_destination:
groups = []
for ids in utils.grouper(content_ids, self.batch_size):
sig_ids = self.task_destination.s(list(ids))
groups.append(sig_ids)
group(groups).delay()
def send_origins(self, origins: Iterable[model.ListedOrigin]) -> int:
"""Record a list of :class:`model.ListedOrigin` in the scheduler.
Returns:
the number of listed origins recorded in the scheduler
"""
count = 0
for batch_origins in grouper(origins, n=1000):
ret = self.scheduler.record_listed_origins(batch_origins)
count += len(ret)
return count
def run(
self,
partition_id: int,
nb_partitions: int,
skip_existing: bool = True,
**kwargs,
) -> Dict:
"""Given a partition of content ids, index the contents within.
Either the indexer is incremental (filter out existing computed data) or it
computes everything from scratch.
Args:
partition_id: Index of the partition to fetch
nb_partitions: Total number of partitions to split into
skip_existing: Skip existing indexed data
(default) or not
**kwargs: passed to the `index` method
Returns:
dict with the indexing task status
"""
summary: Dict[str, Any] = {"status": "uneventful"}
count = 0
try:
if skip_existing:
gen = self._index_with_skipping_already_done(
partition_id, nb_partitions)
else:
gen = self._index_contents(partition_id,
nb_partitions,
indexed=set([]))
count_object_added_key: Optional[str] = None
for contents in utils.grouper(gen,
n=self.config["write_batch_size"]):
res = self.persist_index_computations(list(contents))
if not count_object_added_key:
count_object_added_key = list(res.keys())[0]
count += res[count_object_added_key]
if count > 0:
summary["status"] = "eventful"
except Exception:
if not self.catch_exceptions:
raise
self.log.exception("Problem when computing metadata.")
summary["status"] = "failed"
if count > 0 and count_object_added_key:
summary[count_object_added_key] = count
return summary
def compute_s3_files(db_url, block_size, limit, final, huge):
from swh.scheduler.celery_backend.config import app
from swh.loader.antelink import tasks # noqa
if db_url:
store = storage.Storage(db_url)
gen_data = s3_files_to_download(store, huge, final, limit)
else:
gen_data = utils.gen_path_length_from_stdin()
for paths in grouper(gen_data, block_size):
app.tasks['swh.loader.antelink.tasks.AntelinkS3DownloaderTsk'].delay(
list(paths))
def compute_s3_files(db_url, block_size, limit, final, huge):
from swh.scheduler.celery_backend.config import app
from swh.loader.antelink import tasks # noqa
if db_url:
store = storage.Storage(db_url)
gen_data = s3_files_to_download(store, huge, final, limit)
else:
gen_data = utils.gen_path_length_from_stdin()
for paths in grouper(gen_data, block_size):
app.tasks['swh.loader.antelink.tasks.AntelinkS3DownloaderTsk'].delay(
list(paths))
def mget(self, index_name, doc_ids, chunk_size=500, source=True):
"""Retrieve document's full content according to their ids as per
source's setup.
The `source` allows to retrieve only what's interesting, e.g:
- source=True ; gives back the original indexed data
- source=False ; returns without the original _source field
- source=['task_id'] ; returns only task_id in the _source field
Args:
index_name (str): Name of the concerned index.
doc_ids (generator): Generator of ids to retrieve
chunk_size (int): Number of documents chunk to send for retrieval
source (bool/[str]): Source of information to return
Yields:
document indexed as per source's setup
"""
if isinstance(source, list):
source = {"_source": ",".join(source)}
else:
source = {"_source": str(source).lower()}
for ids in utils.grouper(doc_ids, n=1000):
res = self.storage.mget(
body={"ids": list(ids)},
index=index_name,
doc_type=self.doc_type,
params=source,
)
if not res:
logger.error("Error during retrieval of data, skipping!")
continue
for doc in res["docs"]:
found = doc.get("found")
if not found:
msg = "Doc id %s not found, not indexed yet" % doc["_id"]
logger.warning(msg)
continue
yield doc["_source"]
def get_contents_to_archive(self):
gen_content_ids = (ids for ids in utils.grouper(
read_sha1_from_stdin(), self.config['batch_max_size']))
if self.force_copy:
for content_ids in gen_content_ids:
content_ids = list(content_ids)
if not content_ids:
continue
# Add missing entries in archiver table
self._add_unknown_content_ids(content_ids)
print('Send %s contents to archive' % len(content_ids))
for content_id in content_ids:
# force its status to missing
self.archiver_storage.content_archive_update(
content_id, self.destination, 'missing')
yield content_id
else:
for content_ids in gen_content_ids:
content_ids = list(content_ids)
# Add missing entries in archiver table
self._add_unknown_content_ids(content_ids)
# Filter already copied data
content_ids = list(
self.archiver_storage.content_archive_get_missing(
content_ids=content_ids,
backend_name=self.destination))
if not content_ids:
continue
print('Send %s contents to archive' % len(content_ids))
for content in content_ids:
yield content
def run(self, contents: List[Dict[str, Any]]) -> Dict:
"""Given a list of content:
- (re)compute a given set of checksums on contents available in our
object storage
- update those contents with the new metadata
Args:
contents: contents as dictionary with necessary keys.
key present in such dictionary should be the ones defined in
the 'primary_key' option.
Returns:
A summary dict with key 'status', task' status and 'count' the
number of updated contents.
"""
status = "uneventful"
count = 0
for data in utils.grouper(self.get_new_contents_metadata(contents),
self.batch_size_update):
groups: Dict[str, List[Any]] = defaultdict(list)
for content, keys_to_update in data:
keys_str = ",".join(keys_to_update)
groups[keys_str].append(content)
for keys_to_update, contents in groups.items():
keys: List[str] = keys_to_update.split(",")
try:
self.storage.content_update(contents, keys=keys)
count += len(contents)
status = "eventful"
except Exception:
self.log.exception("Problem during update.")
continue
return {
"status": status,
"count": count,
}
def process_origin_visits(visits, scheduler, task_names):
task_dicts = []
logging.debug("processing origin visits %r", visits)
if task_names.get("origin_metadata"):
visits = [visit for visit in visits if visit["status"] == "full"]
visit_batches = grouper(visits, MAX_ORIGINS_PER_TASK)
for visit_batch in visit_batches:
visit_urls = []
for visit in visit_batch:
if isinstance(visit["origin"], str):
visit_urls.append(visit["origin"])
else:
visit_urls.append(visit["origin"]["url"])
task_dicts.append(
create_task_dict(
task_names["origin_metadata"],
"oneshot",
visit_urls,
retries_left=1,
))
if task_dicts:
scheduler.create_tasks(task_dicts)
def store_file_to_antelink_db_per_block(db, path):
with db.transaction() as cur:
for data in grouper(load_data(path), BLOCK_SIZE):
db.copy_to(data, 'content_s3', ['sha1', 'path', 'length'], cur)
def store_file_to_antelink_db_per_block(db, path):
with db.transaction() as cur:
for data in grouper(load_data(path), BLOCK_SIZE):
db.copy_to(data, 'content_s3',
['sha1', 'path', 'length'], cur)
def flush(
self,
object_types: Sequence[LObjectType] = OBJECT_TYPES
) -> Dict[str, int]:
summary: Dict[str, int] = {}
def update_summary(stats):
for k, v in stats.items():
summary[k] = v + summary.get(k, 0)
for object_type in object_types:
buffer_ = self._objects[object_type]
if not buffer_:
continue
if logger.isEnabledFor(logging.DEBUG):
log = "Flushing %s objects of type %s"
log_args = [len(buffer_), object_type]
if object_type == "content":
log += " (%s bytes)"
log_args.append(
sum(cast(Content, c).length for c in buffer_.values()))
elif object_type == "directory":
log += " (%s entries)"
log_args.append(
sum(
len(cast(Directory, d).entries)
for d in buffer_.values()))
elif object_type == "revision":
log += " (%s parents, %s estimated bytes)"
log_args.extend((
sum(
len(cast(Revision, r).parents)
for r in buffer_.values()),
sum(
estimate_revision_size(cast(Revision, r))
for r in buffer_.values()),
))
elif object_type == "release":
log += " (%s estimated bytes)"
log_args.append(
sum(
estimate_release_size(cast(Release, r))
for r in buffer_.values()))
logger.debug(log, *log_args)
batches = grouper(buffer_.values(),
n=self._buffer_thresholds[object_type])
for batch in batches:
add_fn = getattr(self.storage, "%s_add" % object_type)
stats = add_fn(list(batch))
update_summary(stats)
# Flush underlying storage
stats = self.storage.flush(object_types)
update_summary(stats)
self.clear_buffers(object_types)
return summary
def archive_tasks(
ctx,
before,
after,
batch_index,
bulk_index,
batch_clean,
dry_run,
verbose,
cleanup,
start_from,
):
"""Archive task/task_run whose (task_type is 'oneshot' and task_status
is 'completed') or (task_type is 'recurring' and task_status is
'disabled').
With --dry-run flag set (default), only list those.
"""
from itertools import groupby
from swh.core.utils import grouper
from swh.scheduler.backend_es import ElasticSearchBackend
from swh.scheduler.utils import utcnow
config = ctx.obj["config"]
scheduler = ctx.obj["scheduler"]
if not scheduler:
raise ValueError("Scheduler class (local/remote) must be instantiated")
logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO)
logger = logging.getLogger(__name__)
logging.getLogger("urllib3").setLevel(logging.WARN)
logging.getLogger("elasticsearch").setLevel(logging.ERROR)
if dry_run:
logger.info("**DRY-RUN** (only reading db)")
if not cleanup:
logger.info("**NO CLEANUP**")
es_storage = ElasticSearchBackend(**config)
now = utcnow()
# Default to archive tasks from a rolling month starting the week
# prior to the current one
if not before:
before = now.shift(weeks=-1).format("YYYY-MM-DD")
if not after:
after = now.shift(weeks=-1).shift(months=-1).format("YYYY-MM-DD")
logger.debug(
"index: %s; cleanup: %s; period: [%s ; %s]"
% (not dry_run, not dry_run and cleanup, after, before)
)
def get_index_name(
data: Dict[str, Any], es_storage: ElasticSearchBackend = es_storage
) -> str:
"""Given a data record, determine the index's name through its ending
date. This varies greatly depending on the task_run's
status.
"""
date = data.get("started")
if not date:
date = data["scheduled"]
return es_storage.compute_index_name(date.year, date.month)
def index_data(before, page_token, batch_index):
while True:
result = scheduler.filter_task_to_archive(
after, before, page_token=page_token, limit=batch_index
)
tasks_sorted = sorted(result["tasks"], key=get_index_name)
groups = groupby(tasks_sorted, key=get_index_name)
for index_name, tasks_group in groups:
logger.debug("Index tasks to %s" % index_name)
if dry_run:
for task in tasks_group:
yield task
continue
yield from es_storage.streaming_bulk(
index_name,
tasks_group,
source=["task_id", "task_run_id"],
chunk_size=bulk_index,
)
page_token = result.get("next_page_token")
if page_token is None:
break
gen = index_data(before, page_token=start_from, batch_index=batch_index)
if cleanup:
for task_ids in grouper(gen, n=batch_clean):
task_ids = list(task_ids)
logger.info("Clean up %s tasks: [%s, ...]" % (len(task_ids), task_ids[0]))
if dry_run: # no clean up
continue
ctx.obj["scheduler"].delete_archived_tasks(task_ids)
else:
for task_ids in grouper(gen, n=batch_index):
task_ids = list(task_ids)
logger.info("Indexed %s tasks: [%s, ...]" % (len(task_ids), task_ids[0]))
logger.debug("Done!")
