def test_sparse_query_functionality():
op = SparseFieldsQuery(model=Owner)
assert op.meta()["default_fields"] == [
"name", "age", "weight", "last_updated"
]
assert op.query() == {
"properties": ["name", "age", "weight", "last_updated"]
}
def summary_resource(summary_store):
resource = ReadOnlyResource(
summary_store,
SummaryDoc,
query_operators=[
MaterialIDsSearchQuery(),
FormulaQuery(),
ElementsQuery(),
SymmetryQuery(),
SearchIsStableQuery(),
SearchIsTheoreticalQuery(),
SearchMagneticQuery(),
SearchESQuery(),
NumericQuery(model=SummaryDoc, excluded_fields=["composition"]),
SearchHasReconstructedQuery(),
HasPropsQuery(),
DeprecationQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(SummaryDoc, default_fields=["material_id"]),
],
tags=["Summary"],
disable_validation=True,
)
return resource
async def get_by_key(
key: str = Path(
..., alias=key_name, title=f"The {key_name} of the {model_name} to get"
),
fields: STORE_PARAMS = Depends(
SparseFieldsQuery(self.model, [self.store.key]).query
),
):
f"""
Get's a document by the primary key in the store
Args:
{key_name}: the id of a single {model_name}
Returns:
a single {model_name} document
"""
self.store.connect()
item = self.store.query_one(
criteria={self.store.key: key}, properties=fields["properties"]
)
if item is None:
raise HTTPException(
status_code=404,
detail=f"Item with {self.store.key} = {key} not found",
)
response = {"data": [item]} # , "meta": Meta()}
return response
def xas_resource(xas_store):
resource = ReadOnlyResource(
xas_store,
XASDoc,
query_operators=[
FormulaQuery(),
ElementsQuery(),
XASQuery(),
XASTaskIDQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
XASDoc,
default_fields=[
"xas_id",
"task_id",
"edge",
"absorbing_element",
"formula_pretty",
"spectrum_type",
"last_updated",
],
),
],
tags=["XAS"],
disable_validation=True,
)
return resource
def test_sparse_query_serialization():
op = SparseFieldsQuery(model=Owner)
with ScratchDir("."):
dumpfn(op, "temp.json")
new_op = loadfn("temp.json")
assert new_op.query() == {
"properties": ["name", "age", "weight", "last_updated"]
}
def __init__(
self,
store: Store,
model: Type[BaseModel],
tags: Optional[List[str]] = None,
query_operators: Optional[List[QueryOperator]] = None,
key_fields: Optional[List[str]] = None,
query: Optional[Dict] = None,
enable_get_by_key: bool = True,
enable_default_search: bool = True,
disable_validation: bool = False,
include_in_schema: Optional[bool] = True,
sub_path: Optional[str] = "/",
):
"""
Args:
store: The Maggma Store to get data from
model: The pydantic model this Resource represents
tags: List of tags for the Endpoint
query_operators: Operators for the query language
key_fields: List of fields to always project. Default uses SparseFieldsQuery
to allow user to define these on-the-fly.
enable_get_by_key: Enable default key route for endpoint.
enable_default_search: Enable default endpoint search behavior.
disable_validation: Whether to use ORJSON and provide a direct FastAPI response.
Note this will disable auto JSON serialization and response validation with the
provided model.
include_in_schema: Whether the endpoint should be shown in the documented schema.
sub_path: sub-URL path for the resource.
"""
self.store = store
self.tags = tags or []
self.query = query or {}
self.key_fields = key_fields
self.versioned = False
self.enable_get_by_key = enable_get_by_key
self.enable_default_search = enable_default_search
self.disable_validation = disable_validation
self.include_in_schema = include_in_schema
self.sub_path = sub_path
self.response_model = Response[model] # type: ignore
self.query_operators = (
query_operators if query_operators is not None else [
PaginationQuery(),
SparseFieldsQuery(
model,
default_fields=[
self.store.key, self.store.last_updated_field
],
),
])
super().__init__(model)
def fermi_resource(fermi_store):
resource = ReadOnlyResource(
fermi_store,
FermiDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(FermiDoc, default_fields=["task_id", "last_updated"]),
],
tags=["Electronic Structure"],
disable_validation=True,
)
return resource
def wulff_resource(wulff_store):
resource = ReadOnlyResource(
wulff_store,
WulffDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(WulffDoc, default_fields=["task_id"]),
],
tags=["Surface Properties"],
enable_default_search=False,
disable_validation=True,
)
return resource
def similarity_resource(similarity_store):
resource = ReadOnlyResource(
similarity_store,
SimilarityDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(SimilarityDoc, default_fields=["task_id"]),
],
tags=["Similarity"],
enable_default_search=False,
disable_validation=True,
)
return resource
def dois_resource(dois_store):
resource = ReadOnlyResource(
dois_store,
DOIDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(DOIDoc, default_fields=["task_id", "doi"]),
],
tags=["DOIs"],
enable_default_search=False,
disable_validation=True,
)
return resource
def eos_resource(eos_store):
resource = ReadOnlyResource(
eos_store,
EOSDoc,
query_operators=[
NumericQuery(model=EOSDoc),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(EOSDoc, default_fields=["task_id"]),
],
tags=["EOS"],
disable_validation=True,
)
return resource
def magnetism_resource(magnetism_store):
resource = ReadOnlyResource(
magnetism_store,
MagnetismDoc,
query_operators=[
MagneticQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(MagnetismDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Magnetism"],
disable_validation=True,
)
return resource
def surface_props_resource(surface_prop_store):
resource = ReadOnlyResource(
surface_prop_store,
SurfacePropDoc,
query_operators=[
NumericQuery(model=SurfacePropDoc),
ReconstructedQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(SurfacePropDoc, default_fields=["task_id"]),
],
tags=["Surface Properties"],
disable_validation=True,
)
return resource
def phonon_bsdos_resource(phonon_bs_store):
resource = ReadOnlyResource(
phonon_bs_store,
PhononBSDOSDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(
PhononBSDOSDoc, default_fields=["task_id", "last_updated"]
),
],
tags=["Phonon"],
enable_default_search=False,
disable_validation=True,
)
return resource
def dielectric_resource(dielectric_store):
resource = ReadOnlyResource(
dielectric_store,
DielectricDoc,
query_operators=[
DielectricQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(DielectricDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Dielectric"],
disable_validation=True,
)
return resource
def bs_obj_resource(s3_store):
resource = ReadOnlyResource(
s3_store,
BSObjectDoc,
query_operators=[
ObjectQuery(),
SparseFieldsQuery(BSObjectDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Electronic Structure"],
enable_get_by_key=False,
enable_default_search=True,
sub_path="/bandstructure/object/",
disable_validation=True,
)
return resource
def oxi_states_resource(oxi_states_store):
resource = ReadOnlyResource(
oxi_states_store,
OxidationStateDoc,
query_operators=[
FormulaQuery(),
PossibleOxiStateQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
OxidationStateDoc, default_fields=["material_id", "last_updated"]
),
],
tags=["Oxidation States"],
disable_validation=True,
)
return resource
def molecules_resource(molecules_store):
resource = ReadOnlyResource(
molecules_store,
MoleculesDoc,
query_operators=[
MoleculeBaseQuery(),
MoleculeElementsQuery(),
MoleculeFormulaQuery(),
MultipleTaskIDsQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(MoleculesDoc, default_fields=["task_id"]),
],
tags=["Molecules"],
disable_validation=True,
)
return resource
def __init__(
self,
store: Store,
model: Union[BaseModel, str],
tags: Optional[List[str]] = None,
query_operators: Optional[List[QueryOperator]] = None,
description: str = None,
):
"""
Args:
store: The Maggma Store to get data from
model: the pydantic model to apply to the documents from the Store
This can be a string with a full python path to a model or
an actuall pydantic Model if this is being instantied in python
code. Serializing this via Monty will autoconvert the pydantic model
into a python path string
tags: list of tags for the Endpoint
query_operators: operators for the query language
description: an explanation of wht does this resource do
"""
self.store = store
self.tags = tags or []
self.description = description
self.model: BaseModel = BaseModel()
if isinstance(model, str):
module_path = ".".join(model.split(".")[:-1])
class_name = model.split(".")[-1]
self.model = dynamic_import(module_path, class_name)
else:
self.model = model
self.query_operators = (
query_operators
if query_operators is not None
else [
PaginationQuery(),
SparseFieldsQuery(self.model, default_fields=[self.store.key]),
DefaultDynamicQuery(self.model),
]
)
self.response_model = Response[self.model] # type: ignore
self.router = APIRouter()
self.prepare_endpoint()
def __init__(
self,
store: Store,
model: Type[BaseModel],
tags: Optional[List[str]] = None,
query_operators: Optional[List[QueryOperator]] = None,
key_fields: Optional[List[str]] = None,
query: Optional[Dict] = None,
include_in_schema: Optional[bool] = True,
sub_path: Optional[str] = "/",
):
"""
Args:
store: The Maggma Store to get data from
model: The pydantic model this Resource represents
tags: List of tags for the Endpoint
query_operators: Operators for the query language
key_fields: List of fields to always project. Default uses SparseFieldsQuery
to allow user to define these on-the-fly.
include_in_schema: Whether the endpoint should be shown in the documented schema.
sub_path: sub-URL path for the resource.
"""
self.store = store
self.tags = tags or []
self.query = query or {}
self.key_fields = key_fields
self.versioned = False
self.include_in_schema = include_in_schema
self.sub_path = sub_path
self.response_model = Response[model] # type: ignore
self.query_operators = (
query_operators if query_operators is not None else [
PaginationQuery(),
SparseFieldsQuery(
model,
default_fields=[
self.store.key, self.store.last_updated_field
],
),
])
super().__init__(model)
def es_resource(es_store):
resource = ReadOnlyResource(
es_store,
ElectronicStructureDoc,
query_operators=[
ESSummaryDataQuery(),
FormulaQuery(),
ElementsQuery(),
NumericQuery(model=ElectronicStructureDoc),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(ElectronicStructureDoc,
default_fields=["material_id", "last_updated"]),
],
tags=["Electronic Structure"],
disable_validation=True,
)
return resource
def task_resource(task_store):
resource = ReadOnlyResource(
task_store,
TaskDoc,
query_operators=[
FormulaQuery(),
ElementsQuery(),
MultipleTaskIDsQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
TaskDoc,
default_fields=["task_id", "formula_pretty", "last_updated"],
),
],
tags=["Tasks"],
)
return resource
def charge_density_resource(s3_store):
resource = ReadOnlyResource(
s3_store,
ChgcarDataDoc,
query_operators=[
ChgcarTaskIDQuery(),
PaginationQuery(default_limit=5, max_limit=10),
SparseFieldsQuery(
ChgcarDataDoc,
default_fields=["task_id", "last_updated"],
),
],
tags=["Charge Density"],
enable_default_search=True,
enable_get_by_key=False,
disable_validation=True,
)
return resource
def substrates_resource(substrates_store):
resource = ReadOnlyResource(
substrates_store,
SubstratesDoc,
query_operators=[
SubstrateStructureQuery(),
NumericQuery(model=SubstratesDoc),
StringQueryOperator(
model=SubstratesDoc, excluded_fields=["film_orient", "orient"]
),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(SubstratesDoc, default_fields=["film_id", "sub_id"]),
],
tags=["Substrates"],
enable_get_by_key=False,
disable_validation=True,
)
return resource
def dos_resource(es_store):
resource = ReadOnlyResource(
es_store,
ElectronicStructureDoc,
query_operators=[
DOSDataQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
ElectronicStructureDoc,
default_fields=["material_id", "last_updated", "dos"],
),
],
tags=["Electronic Structure"],
enable_get_by_key=False,
sub_path="/dos/",
disable_validation=True,
)
return resource
def elasticity_resource(elasticity_store):
resource = ReadOnlyResource(
elasticity_store,
ElasticityDoc,
query_operators=[
ChemsysQuery(),
BulkModulusQuery(),
ShearModulusQuery(),
PoissonQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
ElasticityDoc, default_fields=["task_id", "pretty_formula"],
),
],
tags=["Elasticity"],
disable_validation=False,
)
return resource
def insertion_electrodes_resource(insertion_electrodes_store):
resource = ReadOnlyResource(
insertion_electrodes_store,
InsertionElectrodeDoc,
query_operators=[
ElectrodeFormulaQuery(),
WorkingIonQuery(),
ElementsQuery(),
NumericQuery(model=InsertionElectrodeDoc),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
InsertionElectrodeDoc,
default_fields=["battery_id", "last_updated"],
),
],
tags=["Electrodes"],
disable_validation=True,
)
return resource
def search_helper(payload, base: str = "/?", debug=True) -> Response:
"""
Helper function to directly query search endpoints
Args:
store: store f
base: base of the query, default to /query?
client: TestClient generated from FastAPI
payload: query in dictionary format
debug: True = print out the url, false don't print anything
Returns:
request.Response object that contains the response of the correspoding payload
"""
store = MemoryStore("owners", key="name")
store.connect()
store.update([d.dict() for d in owners])
endpoint = ReadOnlyResource(
store,
Owner,
query_operators=[
StringQueryOperator(model=Owner),
NumericQuery(model=Owner),
SparseFieldsQuery(model=Owner),
],
disable_validation=True,
)
app = FastAPI()
app.include_router(endpoint.router)
client = TestClient(app)
print(inspect.signature(NumericQuery(model=Owner).query))
url = base + urlencode(payload)
if debug:
print(url)
res = client.get(url)
json = res.json()
return res, json.get("data", []) # type: ignore
def gb_resource(gb_store):
resource = ReadOnlyResource(
gb_store,
GrainBoundaryDoc,
query_operators=[
GBTaskIDQuery(),
NumericQuery(
model=GrainBoundaryDoc, excluded_fields=["rotation_axis", "gb_plane"]
),
GBStructureQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
GrainBoundaryDoc, default_fields=["task_id", "last_updated"]
),
],
tags=["Grain Boundaries"],
enable_get_by_key=False,
disable_validation=True,
)
return resource
def materials_resource(materials_store):
resource = ReadOnlyResource(
materials_store,
MaterialsDoc,
query_operators=[
FormulaQuery(),
ElementsQuery(),
MultiTaskIDQuery(),
SymmetryQuery(),
DeprecationQuery(),
NumericQuery(model=MaterialsDoc),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
MaterialsDoc,
default_fields=["material_id", "formula_pretty", "last_updated"],
),
],
tags=["Materials"],
disable_validation=True,
)
return resource