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 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 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 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 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 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 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 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
def search_helper(payload,
base: str = "/?",
debug=True) -> Tuple[Response, Any]:
"""
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
"""
owner_store = MemoryStore("owners", key="name")
owner_store.connect()
owner_store.update([d.dict() for d in owners])
pets_store = MemoryStore("pets", key="name")
pets_store.connect()
pets_store.update([jsonable_encoder(d) for d in pets])
resources = {
"owners": [
ReadOnlyResource(
owner_store,
Owner,
query_operators=[
StringQueryOperator(model=Owner), # type: ignore
NumericQuery(model=Owner), # type: ignore
SparseFieldsQuery(model=Owner),
PaginationQuery(),
],
)
],
"pets": [
ReadOnlyResource(
pets_store,
Owner,
query_operators=[
StringQueryOperator(model=Pet),
NumericQuery(model=Pet),
SparseFieldsQuery(model=Pet),
PaginationQuery(),
],
)
],
}
api = API(resources=resources)
client = TestClient(api.app)
url = base + urlencode(payload)
if debug:
print(url)
res = client.get(url)
try:
data = res.json().get("data", [])
except Exception:
data = res.reason
return res, data