elasticsearch¶

def __init__(
    self,
    name: str,
    resource_cls: EntryResource,
    resource_mapper: BaseResourceMapper,
    client: Optional["Elasticsearch"] = None,
):
    """Initialize the ElasticCollection for the given parameters.

    Parameters:
        name: The name of the collection.
        resource_cls: The type of entry resource that is stored by the collection.
        resource_mapper: A resource mapper object that handles aliases and
            format changes between deserialization and response.
        client: A preconfigured Elasticsearch client.

    """
    super().__init__(
        resource_cls=resource_cls,
        resource_mapper=resource_mapper,
        transformer=ElasticTransformer(mapper=resource_mapper),
    )

    self.client = client if client else CLIENT
    self.name = name

    # If we are creating a new collection from scratch, also create the index,
    # otherwise assume it has already been created externally
    if CONFIG.insert_test_data:
        self.create_optimade_index()

def count(self, *args, **kwargs) -> int:
    raise NotImplementedError

@staticmethod
def create_elastic_index_from_mapper(
    resource_mapper: BaseResourceMapper, fields: Iterable[str]
) -> Dict[str, Any]:
    """Create a fallback elastic index based on a resource mapper.

    Arguments:
        resource_mapper: The resource mapper to create the index for.
        fields: The list of fields to use in the index.

    Returns:
        The `body` parameter to pass to `client.indices.create(..., body=...)`.

    """
    return {
        "mappings": {
            "doc": {
                "properties": {
                    resource_mapper.get_optimade_field(field): {"type": "keyword"}
                    for field in fields
                }
            }
        }
    }

def create_optimade_index(self) -> None:
    """Load or create an index that can handle aliased OPTIMADE fields and attach it
    to the current client.

    """
    body = self.predefined_index.get(self.name)
    if body is None:
        body = self.create_elastic_index_from_mapper(
            self.resource_mapper, self.all_fields
        )

    properties = {}
    for field in list(body["mappings"]["doc"]["properties"].keys()):
        properties[self.resource_mapper.get_backend_field(field)] = body[
            "mappings"
        ]["doc"]["properties"].pop(field)
    body["mappings"]["doc"]["properties"] = properties
    self.client.indices.create(index=self.name, body=body, ignore=400)

    LOGGER.debug(f"Created Elastic index for {self.name!r} with body {body}")

def insert(self, data: List[EntryResource]) -> None:
    """Add the given entries to the underlying database.

    Warning:
        No validation is performed on the incoming data.

    Arguments:
        data: The entry resource objects to add to the database.

    """

    def get_id(item):
        if self.name == "links":
            id_ = "%s-%s" % (item["id"], item["type"])
        elif "id" in item:
            id_ = item["id"]
        elif "_id" in item:
            # use the existing MongoDB ids in the test data
            id_ = str(item["_id"])
        else:
            # ES will generate ids
            id_ = None
        item.pop("_id", None)
        return id_

    bulk(
        self.client,
        [
            {
                "_index": self.name,
                "_id": get_id(item),
                "_type": "doc",
                "_source": item,
            }
            for item in data
        ],
    )