Skip to content

Commit

Permalink
retrieve atomate2 tasks from opendata
Browse files Browse the repository at this point in the history
  • Loading branch information
tschaume committed Dec 12, 2024
1 parent 4e47d09 commit 0786122
Showing 1 changed file with 7 additions and 5 deletions.
12 changes: 7 additions & 5 deletions mp_api/client/core/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -476,12 +476,14 @@ def _query_resource(
suffix = suffix.replace("_", "-")

# Paginate over all entries in the bucket.
# This will have to change for when a subset of entries from
# the DB is needed.
is_tasks = "tasks" in suffix
bucket_suffix = "parsed" if is_tasks else "build"
# TODO: change when a subset of entries needed from DB
if "tasks" in suffix:
bucket_suffix, prefix = "parsed", "tasks_atomate2"
else:
bucket_suffix = "build"
prefix = f"collections/{db_version}/{suffix}"

bucket = f"materialsproject-{bucket_suffix}"
prefix = suffix if is_tasks else f"collections/{db_version}/{suffix}"
paginator = self.s3_client.get_paginator("list_objects_v2")
pages = paginator.paginate(Bucket=bucket, Prefix=prefix)

Expand Down

0 comments on commit 0786122

Please sign in to comment.