async def create_one(self,
resource: "EntryResourceCreate") -> "EntryResource":
"""Create a new document in the MongoDB collection based on query parameters.
Update the newly created document with an `"id"` field.
The value will be the string representation of the `"_id"` field.
This will only be done if `"id"` is not already present in `resource`.
Parameters:
resource: The resource to be created.
Returns:
The newly created document as a pydantic model entry resource.
"""
resource.last_modified = datetime.utcnow()
result = await self.collection.insert_one(await clean_python_types(
resource.dict(exclude_unset=True)))
LOGGER.debug(
"Inserted resource %r in DB collection %s with ID %s",
resource,
self.collection.name,
result.inserted_id,
)
if not resource.id:
LOGGER.debug(
"Updating resource with an `id` field equal to str(id_).")
await self.collection.update_one(
{"_id": result.inserted_id},
{"$set": {
"id": str(result.inserted_id)
}})
return self.resource_cls(**self.resource_mapper.map_back(
await self.collection.find_one({"_id": result.inserted_id})))
async def ci_dev_startup() -> None:
"""Function to run at app startup - only relevant for CI or development to add test
data."""
if bool(os.getenv("CI", "")):
LOGGER.info(
"CI detected - Will load test gateways (after dropping the collection)!"
)
elif os.getenv("OPTIMADE_MONGO_DATABASE", "") == "optimade_gateway_dev":
LOGGER.info(
"Running in development mode - Will load test gateways (after dropping the"
"collection)!")
else:
LOGGER.debug("Not in CI or development mode - will start normally.")
return
# Add test gateways
import json
from optimade_gateway.mongo.database import MONGO_DB
from pathlib import Path
test_data = (Path(__file__).parent.parent.joinpath(
".ci/test_gateways.json").resolve())
await MONGO_DB[CONFIG.gateways_collection].drop()
if await MONGO_DB[CONFIG.gateways_collection].count_documents({}) != 0:
raise RuntimeError(
f"Unexpectedly found documents in the {CONFIG.gateways_collection!r} Mongo"
" collection after dropping it ! Found number of documents: "
f"{await MONGO_DB[CONFIG.gateways_collection].count_documents({})}"
)
if not test_data.exists():
raise FileNotFoundError(
f"Could not find test data file with test gateways at {test_data} !"
)
with open(test_data, encoding="utf8") as handle:
data = json.load(handle)
await MONGO_DB[CONFIG.gateways_collection].insert_many(data)
async def load_optimade_providers_databases() -> None: # pylint: disable=too-many-branches,too-many-statements,too-many-locals
"""Load in the providers' OPTIMADE databases from Materials-Consortia
Utilize the Materials-Consortia list of OPTIMADE providers at
[https://providers.optimade.org](https://providers.optimade.org).
Load in all databases with a valid base URL.
"""
import asyncio
import httpx
from optimade import __api_version__
from optimade.models import LinksResponse
from optimade.models.links import LinkType
from optimade.server.routers.utils import BASE_URL_PREFIXES
from optimade_gateway.common.utils import clean_python_types, get_resource_attribute
from optimade_gateway.models.databases import DatabaseCreate
from optimade_gateway.queries.perform import db_get_all_resources
from optimade_gateway.routers.utils import resource_factory
if not CONFIG.load_optimade_providers_databases:
LOGGER.debug(
"Will not load databases from Materials-Consortia list of providers."
)
return
if TYPE_CHECKING or bool(os.getenv("MKDOCS_BUILD",
"")): # pragma: no cover
providers: "Union[httpx.Response, LinksResponse]"
async with httpx.AsyncClient() as client:
providers = await client.get(
f"https://providers.optimade.org/v{__api_version__.split('.', maxsplit=1)[0]}"
"/links")
if providers.is_error:
LOGGER.warning(
"Response from Materials-Consortia's list of OPTIMADE providers was not "
"successful (status code != 200). No databases will therefore be added at "
"server startup.")
return
LOGGER.info(
"Registering Materials-Consortia list of OPTIMADE providers' databases."
)
providers = LinksResponse(**providers.json())
valid_providers = []
for provider in providers.data:
if get_resource_attribute(provider, "id") in ("exmpl", "optimade"):
LOGGER.info(
"- %s (id=%r) - Skipping: Not a real provider.",
get_resource_attribute(provider, "attributes.name", "N/A"),
get_resource_attribute(provider, "id"),
)
continue
if not get_resource_attribute(provider, "attributes.base_url"):
LOGGER.info(
"- %s (id=%r) - Skipping: No base URL information.",
get_resource_attribute(provider, "attributes.name", "N/A"),
get_resource_attribute(provider, "id"),
)
continue
valid_providers.append(provider)
# Run queries to each database using the supported major versioned base URL to get a
# list of the provider's databases.
# There is no need to use ThreadPoolExecutor here, since we want this to block
# everything and then finish, before the server actually starts up.
provider_queries = [
asyncio.create_task(
db_get_all_resources(
database=provider,
endpoint="links",
response_model=LinksResponse,
)) for provider in valid_providers
]
for query in asyncio.as_completed(provider_queries):
provider_databases, provider = await query
LOGGER.info(
"- %s (id=%r) - Processing",
get_resource_attribute(provider, "attributes.name", "N/A"),
get_resource_attribute(provider, "id"),
)
if not provider_databases:
LOGGER.info(" - No OPTIMADE databases found.")
continue
provider_databases = [
db for db in provider_databases if await clean_python_types(
get_resource_attribute(db, "attributes.link_type", "")) ==
LinkType.CHILD.value
]
if not provider_databases:
LOGGER.info(" - No OPTIMADE databases found.")
continue
for database in provider_databases:
if not get_resource_attribute(database, "attributes.base_url"):
LOGGER.info(
" - %s (id=%r) - Skipping: No base URL information.",
get_resource_attribute(database, "attributes.name", "N/A"),
get_resource_attribute(database, "id"),
)
continue
LOGGER.info(
" - %s (id=%r) - Checking versioned base URL and /structures",
get_resource_attribute(database, "attributes.name", "N/A"),
get_resource_attribute(database, "id"),
)
async with httpx.AsyncClient() as client:
try:
db_response = await client.get(
f"{str(get_resource_attribute(database, 'attributes.base_url')).rstrip('/')}" # pylint: disable=line-too-long
f"{BASE_URL_PREFIXES['major']}/structures", )
except httpx.ReadTimeout:
LOGGER.info(
" - %s (id=%r) - Skipping: Timeout while requesting "
"%s/structures.",
get_resource_attribute(database, "attributes.name",
"N/A"),
get_resource_attribute(database, "id"),
BASE_URL_PREFIXES["major"],
)
continue
if db_response.status_code != 200:
LOGGER.info(
" - %s (id=%r) - Skipping: Response from %s/structures is not "
"200 OK.",
get_resource_attribute(database, "attributes.name", "N/A"),
get_resource_attribute(database, "id"),
BASE_URL_PREFIXES["major"],
)
continue
new_id = (f"{get_resource_attribute(provider, 'id')}"
f"/{get_resource_attribute(database, 'id')}" if
len(provider_databases) > 1 else get_resource_attribute(
database, "id"))
registered_database, _ = await resource_factory(
DatabaseCreate(
id=new_id,
**await clean_python_types(
get_resource_attribute(database, "attributes", {})),
))
LOGGER.info(
" - %s (id=%r) - Registered database with id=%r",
get_resource_attribute(database, "attributes.name", "N/A"),
get_resource_attribute(database, "id"),
registered_database.id,
)
async def db_get_all_resources(
database: "Union[LinksResource, Dict[str, Any]]",
endpoint: str,
response_model: "EntryResponseMany",
query_params: str = "",
raw_url: str = None,
) -> "Tuple[List[Union[EntryResource, Dict[str, Any]]], Union[LinksResource, Dict[str, Any]]]": # pylint: disable=line-too-long
"""Recursively retrieve all resources from an entry-listing endpoint
This function keeps pulling the `links.next` link if `meta.more_data_available` is
`True` to ultimately retrieve *all* entries for `endpoint`.
!!! warning
This function can be dangerous if an endpoint with hundreds or thousands of
entries is requested.
Parameters:
database: The OPTIMADE implementation to be queried.
It **must** have a valid base URL and id.
endpoint: The entry-listing endpoint, e.g., `"structures"`.
response_model: The expected OPTIMADE pydantic response model, e.g.,
`optimade.models.StructureResponseMany`.
query_params: URL query parameters to pass to the database.
raw_url: A raw URL to use straight up instead of deriving a URL from `database`,
`endpoint`, and `query_params`.
Returns:
A collected list of successful responses' `data` value and the `database`'s ID.
"""
resulting_resources = []
response, _ = db_find(
database=database,
endpoint=endpoint,
response_model=response_model,
query_params=query_params,
raw_url=raw_url,
)
if isinstance(response, ErrorResponse):
# An errored response will result in no databases from a provider.
LOGGER.error(
"Error while querying database (id=%r). Full response: %s",
get_resource_attribute(database, "id"),
response.json(indent=2),
)
return [], database
resulting_resources.extend(response.data)
if response.meta.more_data_available:
next_page = get_resource_attribute(response, "links.next")
if next_page is None:
LOGGER.error(
"Could not find a 'next' link for an OPTIMADE query request to %r "
"(id=%r). Cannot get all resources from /%s, even though this was asked "
"and `more_data_available` is `True` in the response.",
get_resource_attribute(database, "attributes.name", "N/A"),
get_resource_attribute(database, "id"),
endpoint,
)
return resulting_resources, database
more_resources, _ = await db_get_all_resources(
database=database,
endpoint=endpoint,
response_model=response_model,
query_params=query_params,
raw_url=next_page,
)
resulting_resources.extend(more_resources)
return resulting_resources, database
"""Initialize the MongoDB database."""
from os import getenv
from typing import TYPE_CHECKING
from motor.motor_asyncio import AsyncIOMotorClient
from optimade_gateway.common.config import CONFIG
from optimade_gateway.common.logger import LOGGER
if TYPE_CHECKING or bool(getenv("MKDOCS_BUILD", "")): # pragma: no cover
# pylint: disable=unused-import
from pymongo.database import Database
from pymongo.mongo_client import MongoClient
MONGO_CLIENT: "MongoClient" = AsyncIOMotorClient(
CONFIG.mongo_uri,
appname="optimade-gateway",
readConcernLevel="majority",
readPreference="primary",
w="majority",
)
"""The MongoDB motor client."""
MONGO_DB: "Database" = MONGO_CLIENT[CONFIG.mongo_database]
"""The MongoDB motor database.
This is a representation of the database used for the gateway service."""
LOGGER.info("Using: Real MongoDB (motor) at %s", CONFIG.mongo_uri)
LOGGER.info("Database: %s", CONFIG.mongo_database)
async def post_search(request: Request,
search: Search) -> QueriesResponseSingle:
"""`POST /search`
Coordinate a new OPTIMADE query in multiple databases through a gateway:
1. Search for gateway in DB using `optimade_urls` and `database_ids`
1. Create [`GatewayCreate`][optimade_gateway.models.gateways.GatewayCreate] model
1. `POST` gateway resource to get ID - using functionality of `POST /gateways`
1. Create new [Query][optimade_gateway.models.queries.QueryCreate] resource
1. `POST` Query resource - using functionality of `POST /queries`
1. Return `POST /queries` response -
[`QueriesResponseSingle`][optimade_gateway.models.responses.QueriesResponseSingle]
"""
databases_collection = await collection_factory(CONFIG.databases_collection
)
# NOTE: It may be that the final list of base URLs (`base_urls`) contains the same
# provider(s), but with differring base URLS, if, for example, a versioned base URL
# is supplied.
base_urls = set()
if search.database_ids:
databases = await databases_collection.get_multiple(filter={
"id": {
"$in": await clean_python_types(search.database_ids)
}
})
base_urls |= {
get_resource_attribute(database, "attributes.base_url")
for database in databases if get_resource_attribute(
database, "attributes.base_url") is not None
}
if search.optimade_urls:
base_urls |= {_ for _ in search.optimade_urls if _ is not None}
if not base_urls:
msg = "No (valid) OPTIMADE URLs with:"
if search.database_ids:
msg += (
f"\n Database IDs: {search.database_ids} and corresponding found URLs: "
f"{[get_resource_attribute(database, 'attributes.base_url') for database in databases]}"
)
if search.optimade_urls:
msg += f"\n Passed OPTIMADE URLs: {search.optimade_urls}"
raise BadRequest(detail=msg)
# Ensure all URLs are `pydantic.AnyUrl`s
if not all(isinstance(_, AnyUrl) for _ in base_urls):
raise InternalServerError(
"Could unexpectedly not validate all base URLs as proper URLs.")
databases = await databases_collection.get_multiple(
filter={"base_url": {
"$in": await clean_python_types(base_urls)
}})
if len(databases) == len(base_urls):
# At this point it is expected that the list of databases in `databases`
# is a complete set of databases requested.
gateway = GatewayCreate(databases=databases)
elif len(databases) < len(base_urls):
# There are unregistered databases
current_base_urls = {
get_resource_attribute(database, "attributes.base_url")
for database in databases
}
databases.extend([
LinksResource(
id=(f"{url.user + '@' if url.user else ''}{url.host}"
f"{':' + url.port if url.port else ''}"
f"{url.path.rstrip('/') if url.path else ''}").replace(
".", "__"),
type="links",
attributes=LinksResourceAttributes(
name=(f"{url.user + '@' if url.user else ''}{url.host}"
f"{':' + url.port if url.port else ''}"
f"{url.path.rstrip('/') if url.path else ''}"),
description="",
base_url=url,
link_type=LinkType.CHILD,
homepage=None,
),
) for url in base_urls - current_base_urls
])
else:
LOGGER.error(
"Found more database entries in MongoDB than then number of passed base URLs."
" This suggests ambiguity in the base URLs of databases stored in MongoDB.\n"
" base_urls: %s\n databases %s",
base_urls,
databases,
)
raise InternalServerError(
"Unambiguous base URLs. See logs for more details.")
gateway = GatewayCreate(databases=databases)
gateway, created = await resource_factory(gateway)
if created:
LOGGER.debug("A new gateway was created for a query (id=%r)",
gateway.id)
else:
LOGGER.debug("A gateway was found and reused for a query (id=%r)",
gateway.id)
query = QueryCreate(
endpoint=search.endpoint,
gateway_id=gateway.id,
query_parameters=search.query_parameters,
)
query, created = await resource_factory(query)
if created:
asyncio.create_task(perform_query(url=request.url, query=query))
collection = await collection_factory(CONFIG.queries_collection)
return QueriesResponseSingle(
links=ToplevelLinks(next=None),
data=query,
meta=meta_values(
url=request.url,
data_returned=1,
data_available=await collection.acount(),
more_data_available=False,
**{f"_{CONFIG.provider.prefix}_created": created},
),
)
async def get_search(
request: Request,
response: Response,
search_params: SearchQueryParams = Depends(),
entry_params: EntryListingQueryParams = Depends(),
) -> Union[QueriesResponseSingle, EntryResponseMany, ErrorResponse,
RedirectResponse]:
"""`GET /search`
Coordinate a new OPTIMADE query in multiple databases through a gateway:
1. Create a [`Search`][optimade_gateway.models.search.Search] `POST` data - calling
`POST /search`.
1. Wait [`search_params.timeout`][optimade_gateway.queries.params.SearchQueryParams]
seconds before returning the query, if it has not finished before.
1. Return query - similar to `GET /queries/{query_id}`.
This endpoint works similarly to `GET /queries/{query_id}`, where one passes the query
parameters directly in the URL, instead of first POSTing a query and then going to its
URL. Hence, a
[`QueryResponseSingle`][optimade_gateway.models.responses.QueriesResponseSingle] is
the standard response model for this endpoint.
If the timeout time is reached and the query has not yet finished, the user is
redirected to the specific URL for the query.
If the `as_optimade` query parameter is `True`, the response will be parseable as a
standard OPTIMADE entry listing endpoint like, e.g., `/structures`.
For more information see the
[OPTIMADE specification](https://github.com/Materials-Consortia/OPTIMADE/blob/master/optimade.rst#entry-listing-endpoints).
"""
try:
search = Search(
query_parameters=OptimadeQueryParameters(
**{
field: getattr(entry_params, field)
for field in OptimadeQueryParameters.__fields__
if getattr(entry_params, field)
}),
optimade_urls=search_params.optimade_urls,
endpoint=search_params.endpoint,
database_ids=search_params.database_ids,
)
except ValidationError as exc:
raise BadRequest(detail=(
"A Search object could not be created from the given URL query "
f"parameters. Error(s): {exc.errors}")) from exc
queries_response = await post_search(request, search=search)
if not queries_response.data:
LOGGER.error("QueryResource not found in POST /search response:\n%s",
queries_response)
raise RuntimeError(
"Expected the response from POST /search to return a QueryResource, it did "
"not")
once = True
start_time = time()
while ( # pylint: disable=too-many-nested-blocks
time() < (start_time + search_params.timeout) or once):
# Make sure to run this at least once (e.g., if timeout=0)
once = False
collection = await collection_factory(CONFIG.queries_collection)
query: QueryResource = await collection.get_one(
**{"filter": {
"id": queries_response.data.id
}})
if query.attributes.state == QueryState.FINISHED:
if query.attributes.response and query.attributes.response.errors:
for error in query.attributes.response.errors:
if error.status:
for part in error.status.split(" "):
try:
response.status_code = int(part)
break
except ValueError:
pass
if response.status_code and response.status_code >= 300:
break
else:
response.status_code = 500
if search_params.as_optimade:
return await query.response_as_optimade(url=request.url)
return QueriesResponseSingle(
links=ToplevelLinks(next=None),
data=query,
meta=meta_values(
url=request.url,
data_returned=1,
data_available=await collection.acount(),
more_data_available=False,
),
)
await asyncio.sleep(0.1)
# The query has not yet succeeded and we're past the timeout time -> Redirect to
# /queries/<id>
return RedirectResponse(query.links.self)
async def resource_factory( # pylint: disable=too-many-branches
create_resource: "Union[DatabaseCreate, GatewayCreate, QueryCreate]",
) -> "Tuple[Union[LinksResource, GatewayResource, QueryResource], bool]":
"""Get or create a resource
Currently supported resources:
- `"databases"` ([`DatabaseCreate`][optimade_gateway.models.databases.DatabaseCreate]
->
[`LinksResource`](https://www.optimade.org/optimade-python-tools/api_reference/models/links/#optimade.models.links.LinksResource))
- `"gateways"` ([`GatewayCreate`][optimade_gateway.models.gateways.GatewayCreate] ->
[`GatewayResource`][optimade_gateway.models.gateways.GatewayResource])
- `"queries"` ([`QueryCreate`][optimade_gateway.models.queries.QueryCreate] ->
[`QueryResource`][optimade_gateway.models.queries.QueryResource])
For each of the resources, "uniqueness" is determined in the following way:
=== "Databases"
The `base_url` field is considered unique across all databases.
If a `base_url` is provided via a
[`Link`](https://www.optimade.org/optimade-python-tools/api_reference/models/jsonapi/#optimade.models.jsonapi.Link)
model, the `base_url.href` value is used to query the MongoDB.
=== "Gateways"
The collected list of `databases.attributes.base_url` values is considered unique
across all gateways.
In the database, the search is done as a combination of the length/size of the
`databases`' Python list/MongoDB array and a match on all (using the MongoDB
`$all` operator) of the
[`databases.attributes.base_url`](https://www.optimade.org/optimade-python-tools/api_reference/models/links/#optimade.models.links.LinksResourceAttributes.base_url)
element values, when compared with the `create_resource`.
!!! important
The `database_ids` attribute **must not** contain values that are not also
included in the `databases` attribute, in the form of the IDs for the
individual databases. If this should be the case an
[`OptimadeGatewayError`][optimade_gateway.common.exceptions.OptimadeGatewayError]
will be thrown.
=== "Queries"
The `gateway_id`, `query_parameters`, and `endpoint` fields are collectively
considered to define uniqueness for a
[`QueryResource`][optimade_gateway.models.queries.QueryResource] in the MongoDB
collection.
!!! attention
Only the `/structures` entry endpoint can be queried with multiple expected
responses.
This means the `endpoint` field defaults to `"structures"`, i.e., the
[`StructureResource`](https://www.optimade.org/optimade-python-tools/all_models/#optimade.models.structures.StructureResource)
resource model.
Parameters:
create_resource: The resource to be retrieved or created anew.
Returns:
Two things in a tuple:
- Either a [`GatewayResource`][optimade_gateway.models.gateways.GatewayResource];
a [`QueryResource`][optimade_gateway.models.queries.QueryResource]; or a
[`LinksResource`](https://www.optimade.org/optimade-python-tools/api_reference/models/links/#optimade.models.links.LinksResource)
and
- whether or not the resource was newly created.
"""
created = False
if isinstance(create_resource, DatabaseCreate):
collection_name = CONFIG.databases_collection
base_url = get_resource_attribute(create_resource, "base_url")
mongo_query = {
"$or": [
{"base_url": {"$eq": base_url}},
{"base_url.href": {"$eq": base_url}},
]
}
elif isinstance(create_resource, GatewayCreate):
collection_name = CONFIG.gateways_collection
# One MUST have taken care of database_ids prior to calling `resource_factory()`
database_attr_ids = {_.id for _ in create_resource.databases or []}
unknown_ids = {
database_id
for database_id in create_resource.database_ids or []
if database_id not in database_attr_ids
}
if unknown_ids:
raise OptimadeGatewayError(
"When using `resource_factory()` for `GatewayCreate`, `database_ids` MUST"
f" not include unknown IDs. Passed unknown IDs: {unknown_ids}"
)
mongo_query = {
"databases": {"$size": len(create_resource.databases)},
"databases.attributes.base_url": {
"$all": [_.attributes.base_url for _ in create_resource.databases or []]
},
}
elif isinstance(create_resource, QueryCreate):
collection_name = CONFIG.queries_collection
# Currently only /structures entry endpoints can be queried with multiple
# expected responses.
create_resource.endpoint = (
create_resource.endpoint if create_resource.endpoint else "structures"
)
mongo_query = {
"gateway_id": {"$eq": create_resource.gateway_id},
"query_parameters": {"$eq": create_resource.query_parameters},
"endpoint": {"$eq": create_resource.endpoint},
}
else:
raise TypeError(
"create_resource must be either a GatewayCreate or QueryCreate object not "
f"{type(create_resource)!r}"
)
collection = await collection_factory(collection_name)
result, data_returned, more_data_available, _, _ = await collection.afind(
criteria={"filter": await clean_python_types(mongo_query)}
)
if more_data_available:
raise OptimadeGatewayError(
"more_data_available MUST be False for a single entry response, however it "
f"is {more_data_available}"
)
if result:
if data_returned > 1:
raise OptimadeGatewayError(
f"More than one {result[0].type} were found. IDs of found "
f"{result[0].type}: {[_.id for _ in result]}"
)
if isinstance(result, list):
result = result[0]
else:
if isinstance(create_resource, DatabaseCreate):
# Set required `LinksResourceAttributes` values if not set
if not create_resource.description:
create_resource.description = (
f"{create_resource.name} created by OPTIMADE gateway database "
"registration."
)
if not create_resource.link_type:
create_resource.link_type = LinkType.EXTERNAL
if not create_resource.homepage:
create_resource.homepage = None
elif isinstance(create_resource, GatewayCreate):
# Do not store `database_ids`
if "database_ids" in create_resource.__fields_set__:
create_resource.database_ids = None
create_resource.__fields_set__.remove("database_ids")
elif isinstance(create_resource, QueryCreate):
create_resource.state = QueryState.CREATED
result = await collection.create_one(create_resource)
LOGGER.debug("Created new %s: %r", result.type, result)
created = True
return result, created
async def process_db_response(
response: "Union[ErrorResponse, EntryResponseMany, EntryResponseOne]",
database_id: str,
query: "QueryResource",
gateway: "GatewayResource",
) -> "Union[List[EntryResource], List[Dict[str, Any]], EntryResource, Dict[str, Any], None]": # pylint: disable=line-too-long
"""Process an OPTIMADE database response.
The passed `query` will be updated with the top-level `meta` information:
`data_available`, `data_returned`, and `more_data_available`.
Since, only either `data` or `errors` should ever be present, one or the other will
be either an empty list or `None`.
Parameters:
response: The OPTIMADE database response to be processed.
database_id: The database's `id` under which the returned resources or errors
will be delivered.
query: A resource representing the performed query.
gateway: A resource representing the gateway that was queried.
Returns:
The response's `data`.
"""
results = []
errors = []
LOGGER.debug("Starting to process database_id: %s", database_id)
if isinstance(response, ErrorResponse):
for error in response.errors:
if isinstance(error.id,
str) and error.id.startswith("OPTIMADE_GATEWAY"):
warn(error.detail, OptimadeGatewayWarning)
else:
# The model `ErrorResponse` does not allow the objects in the top-level
# `errors` list to be parsed as dictionaries - they must be a pydantic
# model.
meta_error = {}
if error.meta:
meta_error = error.meta.dict()
meta_error.update({
f"_{CONFIG.provider.prefix}_source_gateway": {
"id": gateway.id,
"type": gateway.type,
"links": {
"self": gateway.links.self
},
},
f"_{CONFIG.provider.prefix}_source_database": {
"id": database_id,
"type": "links",
"links": {
"self":
(str(gateway.links.self).split("gateways",
maxsplit=1)[0] +
f"databases/{database_id}")
},
},
})
error.meta = Meta(**meta_error)
errors.append(error)
data_returned = 0
more_data_available = False
else:
results = response.data
if isinstance(results, list):
data_returned = response.meta.data_returned or len(results)
else:
data_returned = response.meta.data_returned or (0 if not results
else 1)
more_data_available = response.meta.more_data_available or False
data_available = response.meta.data_available or 0
extra_updates = {
"$inc": {
"response.meta.data_available": data_available,
"response.meta.data_returned": data_returned,
}
}
if not get_resource_attribute(
query,
"attributes.response.meta.more_data_available",
False,
disambiguate=False, # Extremely minor speed-up
):
# Keep it True, if set to True once.
extra_updates.update({
"$set": {
"response.meta.more_data_available": more_data_available
}
})
# This ensures an empty list under `response.data.{database_id}` is returned if the
# case is simply that there are no results to return.
if errors:
extra_updates.update(
{"$addToSet": {
"response.errors": {
"$each": errors
}
}})
await update_query(
query,
f"response.data.{database_id}",
results,
operator=None,
**extra_updates,
)
return results
async def update_query( # pylint: disable=too-many-branches
query: "QueryResource",
field: str,
value: "Any",
operator: "Optional[str]" = None,
**mongo_kwargs: "Any",
) -> None:
"""Update a query's `field` attribute with `value`.
If `field` is a dot-separated value, then only the last field part may be a
non-pre-existing field. Otherwise a `KeyError` or `AttributeError` will be raised.
!!! note
This can *only* update a field for a query's `attributes`, i.e., this function
cannot update `id`, `type` or any other top-level resource field.
!!! important
`mongo_kwargs` will not be considered for updating the pydantic model instance.
Parameters:
query: The query to be updated.
field: The `attributes` field (key) to be set.
This can be a dot-separated key value to signify embedded fields.
**Example**: `response.meta`.
value: The (possibly) new value for `field`.
operator: A MongoDB operator to be used for updating `field` with `value`.
**mongo_kwargs: Further MongoDB update filters.
"""
operator = operator or "$set"
if operator and not operator.startswith("$"):
operator = f"${operator}"
update_time = datetime.utcnow()
update_kwargs = {"$set": {"last_modified": update_time}}
if mongo_kwargs:
update_kwargs.update(mongo_kwargs)
if operator and operator == "$set":
update_kwargs["$set"].update({field: value})
elif operator:
if operator in update_kwargs:
update_kwargs[operator].update({field: value})
else:
update_kwargs.update({operator: {field: value}})
# MongoDB
collection = await collection_factory(CONFIG.queries_collection)
result: "UpdateResult" = await collection.collection.update_one(
filter={"id": {
"$eq": query.id
}},
update=await clean_python_types(update_kwargs),
)
if result.matched_count != 1:
LOGGER.error(
("matched_count should have been exactly 1, it was: %s. "
"Returned update_one result: %s"),
result.matched_count,
result.raw_result,
)
# Pydantic model instance
query.attributes.last_modified = update_time
if "." in field:
field_list = field.split(".")
field = getattr(query.attributes, field_list[0])
for field_part in field_list[1:-1]:
if isinstance(field, dict):
field = field.get(field_part)
else:
field = getattr(field, field_part)
if isinstance(field, dict):
field[field_list[-1]] = value
else:
setattr(field, field_list[-1], value)
else:
setattr(query.attributes, field, value)