playground/infrastructure/datastore_client.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Module contains the client to communicate with Google Cloud Datastore
"""
import json
import logging
import os.path
from datetime import datetime
from pathlib import Path
from typing import List

import yaml
from google.cloud import datastore
from tqdm import tqdm

import config
from config import Config, Origin, PrecompiledExample, DatastoreProps
from models import Example, SdkEnum, Dataset, Emulator, ImportFile

from api.v1 import api_pb2


class DatastoreException(Exception):
    pass


# Google Datastore documentation link: https://cloud.google.com/datastore/docs/concepts
class DatastoreClient:
    """DatastoreClient is a datastore client for sending a request to the Google."""

    _datastore_client: datastore.Client

    def __init__(self, project:str, namespace: str):
        self._check_envs()
        self._datastore_client = datastore.Client(
            namespace=namespace, project=project
        )

    def _check_envs(self):
        if Config.SDK_CONFIG is None:
            raise KeyError("SDK_CONFIG environment variable should be specified in os")

    def save_to_cloud_datastore(
        self, examples_from_rep: List[Example], sdk: SdkEnum, origin: Origin
    ):
        """
        Save examples, output and meta to datastore
        Args:
            :param examples_from_rep: examples from the repository for saving to the Cloud Datastore
            :param sdk: sdk from parameters
            :param origin: typed origin const PG_EXAMPLES | TB_EXAMPLES
        """
        # initialise data
        updated_example_ids = set()
        now = datetime.today()

        # retrieve the last schema version
        actual_schema_version_key = self._get_actual_schema_version_key()

        # retrieve all example keys before updating
        examples_ids_before_updating = set(self._get_all_examples(sdk, origin))

        # loop through every example to save them to the Cloud Datastore
        for example in tqdm(examples_from_rep):
            with self._datastore_client.transaction():
                sdk_key = self._get_key(
                    DatastoreProps.SDK_KIND, api_pb2.Sdk.Name(example.sdk)
                )
                example_id = self._make_example_id(origin, sdk, example.tag.name)

                self._datastore_client.put(
                    self._to_example_entity(
                        example, example_id, sdk_key, actual_schema_version_key, origin
                    )
                )

                snippet = self._to_snippet_entity(
                    example, example_id, sdk_key, now, actual_schema_version_key, origin,
                )
                self._datastore_client.put(snippet)

                if not example.tag.always_run:
                    self._datastore_client.put_multi(
                        self._pc_object_entities(example, example_id)
                    )

                self._datastore_client.put(self._to_main_file_entity(example, example_id))
                if example.tag.files:
                    self._datastore_client.put_multi(
                        [
                            self._to_additional_file_entity(example_id, file, idx)
                            for idx, file in enumerate(example.tag.files, start=1)
                        ]
                    )

                if example.tag.datasets:
                    self._datastore_client.put_multi(
                        [
                            self._to_dataset_entity(dataset_id, dataset.file_name)
                            for dataset_id, dataset in example.tag.datasets.items()
                        ]
                    )

                updated_example_ids.add(example_id)

        # delete examples from the Cloud Datastore that are not in the repository
        examples_ids_for_removing = examples_ids_before_updating - updated_example_ids
        logging.info(
            "Start of deleting %d extra playground examples ...",
            len(examples_ids_for_removing),
        )
        for ex_id in examples_ids_for_removing:
            with self._datastore_client.transaction():
                self._datastore_client.delete(
                    self._get_key(DatastoreProps.EXAMPLE_KIND, ex_id)
                )
                self._datastore_client.delete(
                    self._get_key(DatastoreProps.SNIPPET_KIND, ex_id)
                )
                self._datastore_client.delete(self._get_files_key(ex_id, 0))
            pc_objs_keys_for_removing = []
            for example_type in [
                PrecompiledExample.GRAPH_EXTENSION.upper(),
                PrecompiledExample.OUTPUT_EXTENSION.upper(),
                PrecompiledExample.LOG_EXTENSION.upper(),
            ]:
                pc_objs_keys_for_removing.append(
                    self._get_key(
                        DatastoreProps.PRECOMPILED_OBJECT_KIND,
                        f"{ex_id}{config.DatastoreProps.KEY_NAME_DELIMITER}{example_type}",
                    )
                )
            self._datastore_client.delete_multi(pc_objs_keys_for_removing)

        logging.info("Finish of deleting extra playground examples ...")

    def save_catalogs(self):
        """
        Save catalogs to the Cloud Datastore
        """
        # save a schema version entity
        schema_entity = datastore.Entity(
            self._get_key(DatastoreProps.SCHEMA_KIND, "0.0.1"),
            exclude_from_indexes=("descr",),
        )
        schema_entity.update({"descr": "Data initialization: a schema version, SDKs"})
        self._datastore_client.put(schema_entity)

        # save a sdk catalog
        sdk_objs: any = None
        with open(Config.SDK_CONFIG, encoding="utf-8") as sdks:
            sdk_objs = yaml.load(sdks.read(), Loader=yaml.SafeLoader)

        sdk_entities = []
        file_name = Path(Config.SDK_CONFIG).stem
        for key in sdk_objs[file_name]:
            default_example = sdk_objs[file_name][key]["default-example"]
            sdk_entity = datastore.Entity(self._get_key(DatastoreProps.SDK_KIND, key))
            sdk_entity.update({"defaultExample": default_example})
            sdk_entities.append(sdk_entity)

        self._datastore_client.put_multi(sdk_entities)

    def _get_actual_schema_version_key(self) -> datastore.Key:
        schema_names = []
        last_schema_version_query = self._datastore_client.query(
            kind=DatastoreProps.SCHEMA_KIND
        )
        last_schema_version_query.keys_only()
        schema_iterator = last_schema_version_query.fetch()
        schemas = list(schema_iterator)
        if len(schemas) == 0:
            logging.error("Schema versions not found")
            raise DatastoreException(
                "Schema versions not found. Schema versions must be downloaded during application startup"
            )
        for schema in schemas:
            schema_names.append(schema.key.name)
        schema_names.sort(reverse=True)
        return self._get_key(DatastoreProps.SCHEMA_KIND, schema_names[0])

    def _get_all_examples(self, sdk: SdkEnum, origin: Origin) -> List[str]:
        examples_ids_before_updating = []
        all_examples_query = self._datastore_client.query(
            kind=DatastoreProps.EXAMPLE_KIND
        )
        all_examples_query.add_filter(
            "sdk", "=", self._get_key(DatastoreProps.SDK_KIND, api_pb2.Sdk.Name(sdk))
        )
        all_examples_query.add_filter("origin", "=", origin)
        all_examples_query.keys_only()
        examples_iterator = all_examples_query.fetch()
        for example_item in examples_iterator:
            examples_ids_before_updating.append(example_item.key.name)
        return examples_ids_before_updating

    def _get_key(self, kind: str, identifier: str) -> datastore.Key:
        return self._datastore_client.key(kind, identifier)

    def _get_snippet_key(self, snippet_id: str):
        return self._get_key(DatastoreProps.SNIPPET_KIND, snippet_id)

    def _get_example_key(self, example_id: str):
        return self._get_key(DatastoreProps.EXAMPLE_KIND, example_id)

    def _get_dataset_key(self, dataset_id: str):
        return self._get_key(DatastoreProps.DATASET_KIND, dataset_id)

    def _make_example_id(self, origin: Origin, sdk: SdkEnum, name: str):
        # ToB examples (and other related entities: snippets, files, pc_objects)
        # and Beam Documentation examples have origin prefix in a key
        if origin == Origin.TB_EXAMPLES or origin == Origin.PG_BEAMDOC:
            return config.DatastoreProps.KEY_NAME_DELIMITER.join(
                [
                    origin,
                    api_pb2.Sdk.Name(sdk),
                    name,
                ]
            )
        return config.DatastoreProps.KEY_NAME_DELIMITER.join(
            [
                api_pb2.Sdk.Name(sdk),
                name,
            ]
        )

    def _get_files_key(self, example_id: str, idx: int):
        name = config.DatastoreProps.KEY_NAME_DELIMITER.join([example_id, str(idx)])
        return self._get_key(DatastoreProps.FILES_KIND, name)

    def _get_pc_objects_key(self, example_id: str, pc_obj_type: str):
        return self._get_key(
            DatastoreProps.PRECOMPILED_OBJECT_KIND,
            config.DatastoreProps.KEY_NAME_DELIMITER.join([example_id, pc_obj_type]),
        )

    def _to_snippet_entity(
        self,
        example: Example,
        example_id: str,
        sdk_key: datastore.Key,
        now: datetime,
        schema_key: datastore.Key,
        origin: Origin,
    ) -> datastore.Entity:
        snippet_entity = datastore.Entity(self._get_snippet_key(example_id))
        snippet_entity.update(
            {
                "sdk": sdk_key,
                "pipeOpts": self._get_pipeline_options(example),
                "created": now,
                "origin": origin,
                "numberOfFiles": 1 + len(example.tag.files),
                "schVer": schema_key,
                "complexity": f"COMPLEXITY_{example.tag.complexity}",
            }
        )
        if example.tag.datasets:
            snippet_entity.update({"datasets": self._snippet_datasets(example)})
        return snippet_entity

    def _get_pipeline_options(self, example: Example):
        pip_opts = example.tag.pipeline_options
        if pip_opts is not None:
            return pip_opts
        return ""

    def _to_example_entity(
        self,
        example: Example,
        example_id: str,
        sdk_key: datastore.Key,
        schema_key: datastore.Key,
        origin: Origin,
    ) -> datastore.Entity:
        example_entity = datastore.Entity(self._get_example_key(example_id))
        example_entity.update(
            {
                "name": example.tag.name,
                "sdk": sdk_key,
                "descr": example.tag.description,
                "tags": example.tag.tags,
                "cats": example.tag.categories,
                "path": example.url_vcs,  # keep for backward-compatibity, to be removed
                "type": api_pb2.PrecompiledObjectType.Name(example.type),
                "alwaysRun": example.tag.always_run,
                "neverRun": example.tag.never_run,
                "origin": origin,
                "schVer": schema_key,
                "urlVCS": example.url_vcs,
                "urlNotebook": example.tag.url_notebook,
            }
        )
        return example_entity

    def _pc_object_entities(
        self, example: Example, example_id: str
    ) -> List[datastore.Entity]:
        entities = []
        entities.append(
            self._pc_obj_entity(
                example_id,
                example.graph,
                PrecompiledExample.GRAPH_EXTENSION.upper(),
            )
        )
        entities.append(
            self._pc_obj_entity(
                example_id,
                example.output,
                PrecompiledExample.OUTPUT_EXTENSION.upper(),
            )
        )
        entities.append(
            self._pc_obj_entity(
                example_id, example.logs, PrecompiledExample.LOG_EXTENSION.upper()
            )
        )
        return entities

    def _pc_obj_entity(
        self, example_id: str, content: str, pc_obj_type: str
    ) -> datastore.Entity:
        pc_obj_entity = datastore.Entity(
            self._get_pc_objects_key(example_id, pc_obj_type),
            exclude_from_indexes=("content",),
        )
        pc_obj_entity.update({"content": content})
        return pc_obj_entity

    def _to_main_file_entity(self, example: Example, example_id: str):
        file_entity = datastore.Entity(
            self._get_files_key(example_id, 0), exclude_from_indexes=("content",)
        )
        file_entity.update(
            {
                "name": self._get_file_name_with_extension(
                    example.tag.name, example.sdk
                ),
                "content": example.code,
                "cntxLine": example.context_line,
                "isMain": True,
            }
        )
        return file_entity

    def _to_additional_file_entity(self, example_id: str, file: ImportFile, idx: int):
        file_entity = datastore.Entity(
            self._get_files_key(example_id, idx), exclude_from_indexes=("content",)
        )
        file_entity.update(
            {
                "name": file.name,
                "content": file.content,
                "cntxLine": file.context_line,
                "isMain": False,
            }
        )
        return file_entity


    def _to_dataset_entity(self, dataset_id: str, file_name: str):
        dataset_entity = datastore.Entity(self._get_dataset_key(dataset_id))
        dataset_entity.update({"path": file_name})
        return dataset_entity

    def _to_dataset_nested_entity(self, dataset_id: str, emulator: Emulator):
        nested_entity = datastore.Entity()
        nested_entity.update(
            {
                "dataset": self._get_dataset_key(dataset_id),
                "emulator": emulator.type,
                "config": json.dumps({"topic": emulator.topic.id})
            }
        )
        return nested_entity

    def _snippet_datasets(self, example: Example) -> List[datastore.Entity]:
        datasets = []
        for emulator in example.tag.emulators:
            dataset_nested_entity = self._to_dataset_nested_entity(
                emulator.topic.source_dataset, emulator
            )
            datasets.append(dataset_nested_entity)
        return datasets

    def _get_file_name_with_extension(self, name: str, sdk: int) -> str:
        filename, file_extension = os.path.splitext(name)
        if len(file_extension) == 0:
            extension = Config.SDK_TO_EXTENSION[sdk]
            return f"{filename}.{extension}"
        return name