# Copyright 2020 Google LLC
#
# Licensed 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.
"""Utilities for reading pipelines in testdata."""

import os
import shutil

import ml_metadata as mlmd
from ml_metadata.proto import metadata_store_pb2

# constants in a mlmd metadata store instance derived from a tfx pipeline
_TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), 'tfx_oss_0_21')
_TFX_0_21_DB_FILE = 'metadata.sqlite'
_TFX_0_21_PAYLOAD_DIR = \
  '/tmp/tfx-interactive-2020-03-24T21_31_20.888155-mny0kawj'

TFX_0_21_METRICS_ARTIFACT_IDS = (9, 11)
TFX_0_21_MODEL_ARTIFACT_ID = 8
TFX_0_21_MODEL_DATASET_ID = 2
TFX_0_21_MODEL_URI = os.path.join(_TEST_DATA_DIR, 'Trainer/model/6')
TFX_0_21_STATS_ARTIFACT_ID = 3
TFX_0_21_TRAINER_ID = 6


def get_tfx_pipeline_metadata_store(tmp_db_path: str) -> mlmd.MetadataStore:
  """Copies and opens a metadata_store from the testdata tfx pipeline db.

  It migrates the db to the compatible schema at the head. In addition, it
  updates the stored artifacts' uri to the test data db path, so that the test
  code can open the testdata files mentioned in the database.

  Args:
    tmp_db_path: a temp path for copying the pipeline database.

  Returns:
    A ml-metadata store for the copied pipeline db.
  """
  testdata_db_path = os.path.join(_TEST_DATA_DIR, _TFX_0_21_DB_FILE)
  shutil.copyfile(testdata_db_path, tmp_db_path)

  connection_config = metadata_store_pb2.ConnectionConfig(
      sqlite=metadata_store_pb2.SqliteMetadataSourceConfig(
          filename_uri=tmp_db_path,
          connection_mode=metadata_store_pb2.SqliteMetadataSourceConfig.
          READWRITE,
      )
  )
  # The pipeline db is created with mlmd 0.21, the test run from the head
  # may include newer mlmd schema versions. We migrate the db to newer
  # mlmd schema if needed.
  store = mlmd.MetadataStore(connection_config, enable_upgrade_migration=True)
  # The pipeline db is generated with real pipelines in which the payloads of
  # the artifacts are stored in the file system when the pipeline ran. We fix
  # the uri to point to the testdata payloads generated by the pipeline.
  fixed_artifacts = []
  for artifact in store.get_artifacts():
    artifact.uri = artifact.uri.replace(_TFX_0_21_PAYLOAD_DIR, _TEST_DATA_DIR)
    fixed_artifacts.append(artifact)
  store.put_artifacts(fixed_artifacts)
  return store
