├── .bazelrc ├── .github ├── reusable-build │ └── action.yml └── workflows │ ├── build.yml │ ├── cd-docs.yml │ ├── lint.yml │ └── test.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── RELEASE.md ├── WORKSPACE ├── docker-compose.yml ├── docs └── index.md ├── mkdocs.yml ├── ml_metadata ├── .bazelversion ├── BUILD ├── __init__.py ├── errors.py ├── libmysqlclient.BUILD ├── libmysqlclient.WORKSPACE ├── metadata_store │ ├── BUILD │ ├── __init__.py │ ├── conftest.py │ ├── constants.h │ ├── list_operation_query_helper.cc │ ├── list_operation_query_helper.h │ ├── list_operation_query_helper_test.cc │ ├── list_operation_util.cc │ ├── list_operation_util.h │ ├── metadata_access_object.h │ ├── metadata_access_object_factory.cc │ ├── metadata_access_object_factory.h │ ├── metadata_access_object_factory_test.cc │ ├── metadata_access_object_test.cc │ ├── metadata_access_object_test.h │ ├── metadata_source.cc │ ├── metadata_source.h │ ├── metadata_source_test.cc │ ├── metadata_source_test_suite.cc │ ├── metadata_source_test_suite.h │ ├── metadata_store.cc │ ├── metadata_store.go │ ├── metadata_store.h │ ├── metadata_store.py │ ├── metadata_store_factory.cc │ ├── metadata_store_factory.h │ ├── metadata_store_factory_test.cc │ ├── metadata_store_go_swig.cc │ ├── metadata_store_go_wrap.go │ ├── metadata_store_server_main.cc │ ├── metadata_store_service_impl.cc │ ├── metadata_store_service_impl.h │ ├── metadata_store_service_interface.h │ ├── metadata_store_test.cc │ ├── metadata_store_test.go │ ├── metadata_store_test.py │ ├── metadata_store_test_suite.cc │ ├── metadata_store_test_suite.h │ ├── mlmd_types.py │ ├── mlmd_types_test.py │ ├── mysql_metadata_access_object_test.cc │ ├── mysql_metadata_source.cc │ ├── mysql_metadata_source.h │ ├── mysql_metadata_source_test.cc │ ├── mysql_query_config_executor_test.cc │ ├── postgresql_metadata_access_object_test.cc │ ├── postgresql_metadata_source.cc │ ├── postgresql_metadata_source.h │ ├── postgresql_metadata_source_test.cc │ ├── postgresql_query_executor.cc │ ├── postgresql_query_executor.h │ ├── postgresql_query_executor_test.cc │ ├── pywrap │ │ ├── BUILD │ │ ├── __init__.py │ │ └── metadata_store_extension.cc │ ├── query_config_executor.cc │ ├── query_config_executor.h │ ├── query_executor.cc │ ├── query_executor.h │ ├── query_executor_test.cc │ ├── query_executor_test.h │ ├── rdbms_metadata_access_object.cc │ ├── rdbms_metadata_access_object.h │ ├── rdbms_metadata_access_object_test.cc │ ├── rdbms_metadata_access_object_test.h │ ├── simple_types_util.cc │ ├── simple_types_util.h │ ├── simple_types_util_test.cc │ ├── sqlite_metadata_access_object_test.cc │ ├── sqlite_metadata_source.cc │ ├── sqlite_metadata_source.h │ ├── sqlite_metadata_source_test.cc │ ├── sqlite_metadata_source_util.cc │ ├── sqlite_metadata_source_util.h │ ├── sqlite_query_config_executor_test.cc │ ├── sqlite_rdbms_metadata_access_object_test.cc │ ├── test_mysql_metadata_source_initializer.h │ ├── test_postgresql_metadata_source_initializer.h │ ├── test_postgresql_standalone_metadata_source_initializer.cc │ ├── test_standalone_mysql_metadata_source_initializer.cc │ ├── test_util.h │ ├── test_util_test.cc │ ├── transaction_executor.cc │ ├── transaction_executor.h │ ├── transaction_executor_test.cc │ └── types.h ├── ml_metadata.bzl ├── move_generated_files.sh ├── mysql_configure.bzl ├── postgresql.BUILD ├── postgresql.WORKSPACE ├── proto │ ├── BUILD │ ├── __init__.py │ ├── metadata_source.proto │ ├── metadata_store.proto │ ├── metadata_store_service.proto │ └── testing │ │ ├── BUILD │ │ └── mock.proto ├── py.typed ├── query │ ├── BUILD │ ├── filter_query_ast_resolver.cc │ ├── filter_query_ast_resolver.h │ ├── filter_query_ast_resolver_test.cc │ ├── filter_query_builder.cc │ ├── filter_query_builder.h │ └── filter_query_builder_test.cc ├── repo.bzl ├── simple_types │ ├── BUILD │ ├── __init__.py │ ├── proto │ │ ├── BUILD │ │ ├── __init__.py │ │ └── simple_types.proto │ ├── simple_types_constants.cc │ └── simple_types_constants.h ├── third_party │ ├── BUILD │ ├── glog.BUILD │ ├── grpc.patch │ ├── libmysqlclient.patch │ ├── rules_foreign_cc.patch │ ├── six.BUILD │ └── sqlite.BUILD ├── tools │ ├── __init__.py │ ├── dev_debug │ │ └── Dockerfile │ ├── docker_build │ │ ├── Dockerfile.manylinux2010 │ │ └── build_manylinux.sh │ ├── docker_server │ │ ├── Dockerfile │ │ ├── README.md │ │ └── build_docker_image.sh │ └── documentation │ │ └── build_docs.py ├── util │ ├── BUILD │ ├── field_mask_utils.cc │ ├── field_mask_utils.h │ ├── field_mask_utils_test.cc │ ├── metadata_source_query_config.cc │ ├── metadata_source_query_config.h │ ├── metadata_source_query_config_test.cc │ ├── record_parsing_utils.cc │ ├── record_parsing_utils.h │ ├── record_parsing_utils_test.cc │ ├── return_utils.h │ ├── struct_utils.cc │ ├── struct_utils.h │ └── struct_utils_test.cc ├── version.py └── workspace.bzl ├── pyproject.toml ├── requirements-docs.txt ├── setup.py └── test_constraints.txt /.bazelrc: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # This macro is needed in order for mlmd to build with ZetaSQL which can only 16 | # be compiled upon c++17 or higher. 17 | build --cxxopt="-std=c++17" 18 | 19 | # Needed to avoid zetasql proto error. 20 | build --protocopt=--experimental_allow_proto3_optional 21 | 22 | # icu@: In create_linking_context: in call to create_linking_context(), 23 | # parameter 'user_link_flags' is deprecated and will be removed soon. 24 | # It may be temporarily re-enabled by setting --incompatible_require_linker_input_cc_api=false 25 | build --incompatible_require_linker_input_cc_api=false 26 | build:macos --apple_platform_type=macos 27 | build:macos_arm64 --cpu=darwin_arm64 28 | -------------------------------------------------------------------------------- /.github/reusable-build/action.yml: -------------------------------------------------------------------------------- 1 | name: Resusable steps to build ml-metadata 2 | 3 | inputs: 4 | python-version: 5 | description: 'Python version' 6 | required: true 7 | runs: 8 | using: 'composite' 9 | steps: 10 | 11 | - name: Set up Python ${{ inputs.python-version }} 12 | uses: actions/setup-python@v5 13 | with: 14 | python-version: ${{ inputs.python-version }} 15 | 16 | - name: Upgrade pip 17 | shell: bash 18 | run: | 19 | python -m pip install --upgrade pip pytest 20 | 21 | - name: Build the package for Python ${{ inputs.python-version }} 22 | shell: bash 23 | run: | 24 | version="${{ inputs.python-version }}" 25 | DOCKER_SERVICE=manylinux-python$(echo "$version" | sed 's/\.//') 26 | docker compose build ${DOCKER_SERVICE} 27 | docker compose run ${DOCKER_SERVICE} 28 | 29 | - name: Upload wheel artifact for Python ${{ inputs.python-version }} 30 | uses: actions/upload-artifact@v4.4.0 31 | with: 32 | name: ml-metadata-wheel-py${{ inputs.python-version }} 33 | path: dist/*.whl 34 | 35 | - name: Install built wheel 36 | shell: bash 37 | run: pip install dist/*.whl 38 | -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: Build ml-metadata 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | pull_request: 8 | branches: 9 | - master 10 | workflow_dispatch: 11 | 12 | jobs: 13 | build: 14 | runs-on: ubuntu-latest 15 | strategy: 16 | matrix: 17 | python-version: ["3.9", "3.10", "3.11"] 18 | 19 | steps: 20 | - name: Checkout 21 | uses: actions/checkout@v4 22 | 23 | - name: Build ml-metadata 24 | id: build-ml-metadata 25 | uses: ./.github/reusable-build 26 | with: 27 | python-version: ${{ matrix.python-version }} 28 | 29 | upload_to_pypi: 30 | name: Upload to PyPI 31 | runs-on: ubuntu-latest 32 | if: (github.event_name == 'release' && startsWith(github.ref, 'refs/tags')) || (github.event_name == 'workflow_dispatch') 33 | needs: [build] 34 | environment: 35 | name: pypi 36 | url: https://pypi.org/p/ml-metadata/ 37 | permissions: 38 | id-token: write 39 | steps: 40 | - name: Retrieve wheels 41 | uses: actions/download-artifact@v4.1.8 42 | with: 43 | merge-multiple: true 44 | path: wheels 45 | 46 | - name: List the build artifacts 47 | run: | 48 | ls -lAs wheels/ 49 | 50 | - name: Upload to PyPI 51 | uses: pypa/gh-action-pypi-publish@release/v1.9 52 | with: 53 | packages_dir: wheels/ 54 | -------------------------------------------------------------------------------- /.github/workflows/cd-docs.yml: -------------------------------------------------------------------------------- 1 | name: deploy-docs 2 | on: 3 | workflow_dispatch: 4 | push: 5 | branches: 6 | - 'master' 7 | pull_request: 8 | permissions: 9 | contents: write 10 | jobs: 11 | deploy: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout repo 15 | uses: actions/checkout@v4 16 | 17 | - name: Set up Python 3.12 18 | uses: actions/setup-python@v5 19 | with: 20 | python-version: '3.12' 21 | cache: 'pip' 22 | cache-dependency-path: | 23 | setup.py 24 | 25 | - name: Save time for cache for mkdocs 26 | run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 27 | 28 | - name: Caching 29 | uses: actions/cache@v4 30 | with: 31 | key: mkdocs-material-${{ env.cache_id }} 32 | path: .cache 33 | restore-keys: | 34 | mkdocs-material- 35 | 36 | - name: Install Dependencies 37 | run: pip install -r requirements-docs.txt 38 | 39 | - name: Deploy to GitHub Pages 40 | if: (github.event_name != 'pull_request') 41 | run: | 42 | git config user.name github-actions[bot] 43 | git config user.email 41898282+github-actions[bot]@users.noreply.github.com 44 | mkdocs gh-deploy --force 45 | 46 | - name: Build docs to check for errors 47 | run: mkdocs build 48 | if: (github.event_name == 'pull_request') 49 | -------------------------------------------------------------------------------- /.github/workflows/lint.yml: -------------------------------------------------------------------------------- 1 | name: pre-commit 2 | 3 | on: 4 | pull_request: 5 | push: 6 | branches: [master] 7 | 8 | jobs: 9 | pre-commit: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v4 13 | - uses: actions/setup-python@v5 14 | - uses: pre-commit/action@v3.0.1 15 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test ml-metadata 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | pull_request: 8 | branches: 9 | - master 10 | workflow_dispatch: 11 | 12 | jobs: 13 | test: 14 | runs-on: ubuntu-latest 15 | strategy: 16 | matrix: 17 | python-version: ["3.9", "3.10", "3.11"] 18 | 19 | steps: 20 | - name: Checkout 21 | uses: actions/checkout@v4 22 | 23 | - name: Build ml-metadata 24 | id: build-ml-metadata 25 | uses: ./.github/reusable-build 26 | with: 27 | python-version: ${{ matrix.python-version }} 28 | 29 | - name: Test 30 | run: | 31 | # cleanup (interferes with tests) 32 | rm -rf bazel-* 33 | # run tests 34 | pytest -vv 35 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # LINT.IfChange 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | **/*.py[cod] 5 | **/*$py.class 6 | 7 | # C extensions 8 | **/*.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | pip-wheel-metadata/ 25 | share/python-wheels/ 26 | *.egg-info/ 27 | .installed.cfg 28 | *.egg 29 | MANIFEST 30 | 31 | # PyInstaller 32 | # Usually these files are written by a python script from a template 33 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 34 | *.manifest 35 | *.spec 36 | 37 | # Installer logs 38 | pip-log.txt 39 | pip-delete-this-directory.txt 40 | 41 | # Unit test / coverage reports 42 | htmlcov/ 43 | .tox/ 44 | .nox/ 45 | .coverage 46 | .coverage.* 47 | .cache 48 | nosetests.xml 49 | coverage.xml 50 | *.cover 51 | *.py,cover 52 | .hypothesis/ 53 | .pytest_cache/ 54 | cover/ 55 | 56 | # Translations 57 | *.mo 58 | *.pot 59 | 60 | # Django stuff: 61 | *.log 62 | local_settings.py 63 | db.sqlite3 64 | db.sqlite3-journal 65 | 66 | # Flask stuff: 67 | instance/ 68 | .webassets-cache 69 | 70 | # Scrapy stuff: 71 | .scrapy 72 | 73 | # Sphinx documentation 74 | docs/_build/ 75 | 76 | # PyBuilder 77 | .pybuilder/ 78 | target/ 79 | 80 | # Jupyter Notebook 81 | .ipynb_checkpoints 82 | 83 | # IPython 84 | profile_default/ 85 | ipython_config.py 86 | 87 | # pyenv 88 | .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 110 | .pdm.toml 111 | .pdm-python 112 | .pdm-build/ 113 | 114 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 115 | __pypackages__/ 116 | 117 | # Celery stuff 118 | celerybeat-schedule 119 | celerybeat.pid 120 | 121 | # SageMath parsed files 122 | *.sage.py 123 | 124 | # Environments 125 | .env 126 | .venv 127 | env/ 128 | venv/ 129 | ENV/ 130 | env.bak/ 131 | venv.bak/ 132 | 133 | # Spyder project settings 134 | .spyderproject 135 | .spyproject 136 | 137 | # Rope project settings 138 | .ropeproject 139 | 140 | # Intellij project settings 141 | .idea 142 | 143 | # VSCode project settings 144 | .vscode 145 | 146 | # mkdocs documentation 147 | /site 148 | 149 | # mypy 150 | .mypy_cache/ 151 | .dmypy.json 152 | dmypy.json 153 | 154 | # Pyre type checker 155 | .pyre/ 156 | 157 | # Bazel generated files 158 | bazel-* 159 | **/*_pb2.py 160 | **/*_pb2_grpc.py 161 | # LINT.ThenChange(.dockerignore) 162 | 163 | MODULE.bazel 164 | MODULE.bazel.lock 165 | # pytype static type analyzer 166 | .pytype/ 167 | 168 | # Cython debug symbols 169 | cython_debug/ 170 | 171 | # PyCharm 172 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 173 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 174 | # and can be added to the global gitignore or merged into this file. For a more nuclear 175 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 176 | #.idea/ 177 | 178 | # Bazel build files 179 | bazel-* 180 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # pre-commit is a tool to perform a predefined set of tasks manually and/or 2 | # automatically before git commits are made. 3 | # 4 | # Config reference: https://pre-commit.com/#pre-commit-configyaml---top-level 5 | # 6 | # Common tasks 7 | # 8 | # - Register git hooks: pre-commit install 9 | # - Run on all files: pre-commit run --all-files 10 | # 11 | # These pre-commit hooks are run as CI. 12 | # 13 | # NOTE: if it can be avoided, add configs/args in pyproject.toml or below instead of creating a new `.config.file`. 14 | # https://pre-commit.ci/#configuration 15 | ci: 16 | autoupdate_schedule: monthly 17 | autofix_commit_msg: | 18 | [pre-commit.ci] Apply automatic pre-commit fixes 19 | 20 | repos: 21 | - repo: https://github.com/pre-commit/pre-commit-hooks 22 | rev: v4.6.0 23 | hooks: 24 | - id: end-of-file-fixer 25 | exclude: '\.svg$' 26 | - id: trailing-whitespace 27 | exclude: '\.svg$' 28 | - id: check-json 29 | - id: check-yaml 30 | args: [--allow-multiple-documents, --unsafe] 31 | - id: check-toml 32 | 33 | - repo: https://github.com/astral-sh/ruff-pre-commit 34 | rev: v0.5.6 35 | hooks: 36 | - id: ruff 37 | args: ["--fix"] 38 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | We'd love to accept your patches and contributions to this project. There are 4 | just a few small guidelines you need to follow. 5 | 6 | ## Contributor License Agreement 7 | 8 | Contributions to this project must be accompanied by a Contributor License 9 | Agreement. You (or your employer) retain the copyright to your contribution, 10 | this simply gives us permission to use and redistribute your contributions as 11 | part of the project. Head over to to see 12 | your current agreements on file or to sign a new one. 13 | 14 | You generally only need to submit a CLA once, so if you've already submitted one 15 | (even if it was for a different project), you probably don't need to do it 16 | again. 17 | 18 | ## Code reviews 19 | 20 | All submissions, including submissions by project members, require review. We 21 | use GitHub pull requests for this purpose. Consult 22 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 23 | information on using pull requests. 24 | 25 | ## Using `pre-commit` 26 | 27 | To enforce style and coding standards, this repository uses 28 | [`pre-commit`](https://pre-commit.com/). To get started: 29 | 30 | ```shell 31 | pip install pre-commit 32 | pre-commit install 33 | ``` 34 | 35 | ### Run pre-commit on all files 36 | 37 | By default, `pre-commit` will run its checks on files that have been modified in 38 | a commit. To instead run it on all files, use this command: 39 | 40 | ```console 41 | $ pre-commit run --all-files 42 | 43 | # Alternatively 44 | $ pre-commit run -a 45 | ``` 46 | 47 | ### Skip the pre-commit checks 48 | 49 | Run the following command: 50 | 51 | ```console 52 | $ git commit --no-verify 53 | ``` 54 | 55 | Be aware the pre-commit checks are run on CI, so any violations will need to be 56 | fixed before code gets merged. 57 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # ML Metadata 3 | 4 | [![Python](https://img.shields.io/badge/python%7C3.9%7C3.10%7C3.11-blue)](https://github.com/google/ml-metadata) 5 | [![PyPI](https://badge.fury.io/py/ml-metadata.svg)](https://badge.fury.io/py/ml-metadata) 6 | 7 | *ML Metadata (MLMD)* is a library for recording and retrieving metadata 8 | associated with ML developer and data scientist workflows. 9 | 10 | NOTE: ML Metadata may be backwards incompatible before version 1.0. 11 | 12 | ## Getting Started 13 | 14 | For more background on MLMD and instructions on using it, see the 15 | [getting started guide](https://github.com/google/ml-metadata/blob/master/g3doc/get_started.md) 16 | 17 | ## Installing from PyPI 18 | 19 | The recommended way to install ML Metadata is to use the 20 | [PyPI package](https://pypi.org/project/ml-metadata/): 21 | 22 | ```bash 23 | pip install ml-metadata 24 | ``` 25 | 26 | Then import the relevant packages: 27 | 28 | ```python 29 | from ml_metadata import metadata_store 30 | from ml_metadata.proto import metadata_store_pb2 31 | ``` 32 | 33 | ### Nightly Packages 34 | 35 | ML Metadata (MLMD) also hosts nightly packages at 36 | https://pypi-nightly.tensorflow.org on Google Cloud. To install the latest 37 | nightly package, please use the following command: 38 | 39 | ```bash 40 | pip install --extra-index-url https://pypi-nightly.tensorflow.org/simple ml-metadata 41 | ``` 42 | 43 | ## Installing with Docker 44 | 45 | This is the recommended way to build ML Metadata under Linux, and is 46 | continuously tested at Google. 47 | 48 | Please first install `docker` and `docker-compose` by following the directions: 49 | [docker](https://docs.docker.com/install/); 50 | [docker-compose](https://docs.docker.com/compose/install/). 51 | 52 | Then, run the following at the project root: 53 | 54 | ```bash 55 | DOCKER_SERVICE=manylinux-python${PY_VERSION} 56 | sudo docker compose build ${DOCKER_SERVICE} 57 | sudo docker compose run ${DOCKER_SERVICE} 58 | ``` 59 | 60 | where `PY_VERSION` is one of `{39, 310, 311}`. 61 | 62 | A wheel will be produced under `dist/`, and installed as follows: 63 | 64 | ```shell 65 | pip install dist/*.whl 66 | ``` 67 | 68 | ## Installing from source 69 | 70 | 71 | ### 1. Prerequisites 72 | 73 | To compile and use ML Metadata, you need to set up some prerequisites. 74 | 75 | 76 | #### Install Bazel 77 | 78 | If Bazel is not installed on your system, install it now by following [these 79 | directions](https://bazel.build/versions/master/docs/install.html). 80 | 81 | #### Install cmake 82 | If cmake is not installed on your system, install it now by following [these 83 | directions](https://cmake.org/install/). 84 | 85 | ### 2. Clone ML Metadata repository 86 | 87 | ```shell 88 | git clone https://github.com/google/ml-metadata 89 | cd ml-metadata 90 | ``` 91 | 92 | Note that these instructions will install the latest master branch of ML 93 | Metadata. If you want to install a specific branch (such as a release branch), 94 | pass `-b ` to the `git clone` command. 95 | 96 | ### 3. Build the pip package 97 | 98 | ML Metadata uses Bazel to build the pip package from source: 99 | 100 | ```shell 101 | python setup.py bdist_wheel 102 | ``` 103 | 104 | You can find the generated `.whl` file in the `dist` subdirectory. 105 | 106 | ### 4. Install the pip package 107 | 108 | ```shell 109 | pip install dist/*.whl 110 | ``` 111 | 112 | ### 5.(Optional) Build the grpc server 113 | 114 | ML Metadata uses Bazel to build the c++ binary from source: 115 | 116 | ```shell 117 | bazel build -c opt --define grpc_no_ares=true //ml_metadata/metadata_store:metadata_store_server 118 | ``` 119 | 120 | ## Supported platforms 121 | 122 | MLMD is built and tested on the following 64-bit operating systems: 123 | 124 | * macOS 10.14.6 (Mojave) or later. 125 | * Ubuntu 20.04 or later. 126 | * [DEPRECATED] Windows 10 or later. For a Windows-compatible library, please 127 | refer to MLMD 1.14.0 or earlier versions. 128 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | version: '3.1' 16 | 17 | # Extensions are not support until 3.4 thus the repeated boilerplate below. 18 | 19 | # We mount the MLMD project root at /build (which is the WORKDIR of the image) 20 | # in the container. 21 | services: 22 | manylinux-python39: 23 | image: mlmd-build:manylinux2010 24 | build: 25 | context: . 26 | dockerfile: ml_metadata/tools/docker_build/Dockerfile.manylinux2010 27 | environment: 28 | PYTHON_VERSION: 39 29 | volumes: 30 | - .:/build:delegated 31 | manylinux-python310: 32 | image: mlmd-build:manylinux2010 33 | build: 34 | context: . 35 | dockerfile: ml_metadata/tools/docker_build/Dockerfile.manylinux2010 36 | environment: 37 | PYTHON_VERSION: 310 38 | volumes: 39 | - .:/build:delegated 40 | manylinux-python311: 41 | image: mlmd-build:manylinux2010 42 | build: 43 | context: . 44 | dockerfile: ml_metadata/tools/docker_build/Dockerfile.manylinux2010 45 | environment: 46 | PYTHON_VERSION: 311 47 | volumes: 48 | - .:/build:delegated 49 | -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: "ML Metadata" 2 | repo_name: "ML Metadata" 3 | repo_url: https://github.com/google/ml-metadata 4 | 5 | theme: 6 | name: material 7 | palette: 8 | # Palette toggle for automatic mode 9 | - media: "(prefers-color-scheme)" 10 | primary: custom 11 | accent: custom 12 | toggle: 13 | icon: material/brightness-auto 14 | name: Switch to light mode 15 | 16 | # Palette toggle for light mode 17 | - media: "(prefers-color-scheme: light)" 18 | primary: custom 19 | accent: custom 20 | scheme: default 21 | toggle: 22 | icon: material/brightness-7 23 | name: Switch to dark mode 24 | 25 | # Palette toggle for dark mode 26 | - media: "(prefers-color-scheme: dark)" 27 | primary: custom 28 | accent: custom 29 | scheme: slate 30 | toggle: 31 | icon: material/brightness-4 32 | name: Switch to system preference 33 | favicon: images/favicon.png 34 | 35 | features: 36 | - content.code.copy 37 | - content.code.select 38 | plugins: 39 | - search 40 | - autorefs 41 | - mkdocstrings: 42 | default_handler: python 43 | handlers: 44 | python: 45 | options: 46 | show_source: true 47 | show_root_heading: true 48 | unwrap_annotated: true 49 | show_symbol_type_toc: true 50 | show_symbol_type_heading: true 51 | merge_init_into_class: true 52 | show_signature_annotations: true 53 | separate_signature: true 54 | signature_crossrefs: true 55 | group_by_category: true 56 | show_category_heading: true 57 | inherited_members: true 58 | show_submodules: true 59 | show_object_full_path: false 60 | show_root_full_path: true 61 | docstring_section_style: "spacy" 62 | summary: true 63 | filters: 64 | - "!^_" 65 | - "^__init__$" 66 | - "^__call__$" 67 | - "!^logger" 68 | - "!_test$" 69 | extensions: 70 | - griffe_inherited_docstrings 71 | import: 72 | - https://docs.python.org/3/objects.inv 73 | - caption: 74 | figure: 75 | ignore_alt: true 76 | 77 | markdown_extensions: 78 | - admonition 79 | - attr_list 80 | - toc: 81 | permalink: true 82 | - pymdownx.highlight: 83 | anchor_linenums: true 84 | linenums: false 85 | line_spans: __span 86 | pygments_lang_class: true 87 | - pymdownx.inlinehilite 88 | - pymdownx.snippets 89 | - pymdownx.superfences 90 | - pymdownx.arithmatex: 91 | generic: true 92 | - md_in_html 93 | - pymdownx.emoji: 94 | emoji_index: !!python/name:material.extensions.emoji.twemoji 95 | emoji_generator: !!python/name:material.extensions.emoji.to_svg 96 | 97 | extra_css: 98 | - stylesheets/extra.css 99 | 100 | extra_javascript: 101 | - javascripts/mathjax.js 102 | - https://unpkg.com/mathjax@3/es5/tex-mml-chtml.js 103 | 104 | watch: 105 | - ml_metadata 106 | nav: 107 | - Guide: index.md 108 | - Tutorial: https://tensorflow.github.io/tfx/tutorials/mlmd/mlmd_tutorial/ 109 | - API: 110 | - mlmd: 111 | - Overview: api/mlmd/index.md 112 | - mlmd: api/mlmd/mlmd.md 113 | - mlmd.errors: 114 | - Overview: api/mlmd.errors/index.md 115 | - mlmd.errors: api/mlmd.errors/mlmd.errors.md 116 | - mlmd.proto: 117 | - Overview: api/mlmd.proto/index.md 118 | - mlmd.proto: api/mlmd.proto/mlmd.proto.md 119 | -------------------------------------------------------------------------------- /ml_metadata/.bazelversion: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Google LLC. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 6.1.0 15 | -------------------------------------------------------------------------------- /ml_metadata/BUILD: -------------------------------------------------------------------------------- 1 | load("@bazel_skylib//lib:selects.bzl", "selects") 2 | 3 | licenses(["notice"]) # Apache 2.0 4 | 5 | package(default_visibility = ["//visibility:public"]) 6 | 7 | config_setting( 8 | name = "macos_x86_64", 9 | values = { 10 | "apple_platform_type": "macos", 11 | "cpu": "darwin", 12 | }, 13 | ) 14 | 15 | config_setting( 16 | name = "macos_arm64", 17 | values = { 18 | "apple_platform_type": "macos", 19 | "cpu": "darwin_arm64", 20 | }, 21 | ) 22 | 23 | selects.config_setting_group( 24 | name = "macos", 25 | match_any = [ 26 | ":macos_x86_64", 27 | ":macos_arm64", 28 | ], 29 | ) 30 | 31 | _public_protos = [ 32 | "//ml_metadata/simple_types/proto:simple_types_pb2.py", 33 | "//ml_metadata/proto:metadata_store_pb2.py", 34 | "//ml_metadata/proto:metadata_store_service_pb2.py", 35 | "//ml_metadata/proto:metadata_store_service_pb2_grpc.py", 36 | ] 37 | 38 | _py_extension = select({ 39 | "//conditions:default": [ 40 | "//ml_metadata/metadata_store/pywrap:metadata_store_extension.so", 41 | ], 42 | }) 43 | 44 | sh_binary( 45 | name = "move_generated_files", 46 | srcs = ["move_generated_files.sh"], 47 | data = _py_extension + _public_protos, 48 | ) 49 | -------------------------------------------------------------------------------- /ml_metadata/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Init module for ML Metadata.""" 15 | 16 | # pylint: disable=g-import-not-at-top 17 | try: 18 | from ml_metadata import proto 19 | 20 | # Import metadata_store API. 21 | from ml_metadata.metadata_store import ( 22 | ListOptions, 23 | MetadataStore, 24 | OrderByField, 25 | downgrade_schema, 26 | ) 27 | 28 | # Import version string. 29 | from ml_metadata.version import __version__ 30 | 31 | except ImportError as err: 32 | import sys # pylint: enable=g-import-not-at-top 33 | 34 | sys.stderr.write(f"Error importing: {err}") 35 | 36 | 37 | __all__ = [ 38 | "downgrade_schema", 39 | "ListOptions", 40 | "MetadataStore", 41 | "OrderByField", 42 | ] 43 | -------------------------------------------------------------------------------- /ml_metadata/libmysqlclient.BUILD: -------------------------------------------------------------------------------- 1 | # Source files generated with cmake. 2 | configure_out_srcs = [ 3 | "build/libmariadb/ma_client_plugin.c", 4 | ] 5 | 6 | # Header files generated with cmake. 7 | configure_out_hdrs = [ 8 | "build/include/ma_config.h", 9 | "build/include/config.h", 10 | "build/include/mariadb_version.h", 11 | ] 12 | 13 | # A genrule to run cmake and generate configure_out_(srcs,hdrs). 14 | genrule( 15 | name = "configure", 16 | srcs = glob( 17 | ["**"], 18 | exclude = ["bazel*/**"], 19 | ), 20 | outs = configure_out_srcs + configure_out_hdrs, 21 | cmd = "\n".join([ 22 | "export INSTALL_DIR=$$(pwd)/$(@D)/build", 23 | "export TMP_DIR=$$(mktemp -d -t build.XXXXXX)", 24 | "mkdir -p $$TMP_DIR", 25 | "cp -R $$(pwd)/external/libmysqlclient/* $$TMP_DIR", 26 | "cd $$TMP_DIR", 27 | "mkdir build", 28 | "cd build", 29 | "cmake .. -DCMAKE_BUILD_TYPE=Release $${CMAKE_ICONV_FLAG-}", 30 | "cd ..", 31 | "cp -R ./build/* $$INSTALL_DIR", 32 | "rm -rf $$TMP_DIR", 33 | ]), 34 | ) 35 | 36 | config_setting( 37 | name = "darwin", 38 | values = {"cpu": "darwin"}, 39 | visibility = ["//visibility:public"], 40 | ) 41 | 42 | cc_library( 43 | name = "libmysqlclient", 44 | srcs = configure_out_srcs + [ 45 | # plugins. 46 | "plugins/auth/my_auth.c", 47 | "plugins/auth/old_password.c", 48 | "plugins/pvio/pvio_socket.c", 49 | # ssl. 50 | "libmariadb/secure/openssl.c", 51 | # core libmariadb 52 | "libmariadb/ma_array.c", 53 | "libmariadb/ma_charset.c", 54 | "libmariadb/ma_hash.c", 55 | "libmariadb/ma_net.c", 56 | "libmariadb/mariadb_charset.c", 57 | "libmariadb/ma_time.c", 58 | "libmariadb/ma_default.c", 59 | "libmariadb/ma_errmsg.c", 60 | "libmariadb/mariadb_lib.c", 61 | "libmariadb/ma_list.c", 62 | "libmariadb/ma_pvio.c", 63 | "libmariadb/ma_tls.c", 64 | "libmariadb/ma_alloc.c", 65 | "libmariadb/ma_compress.c", 66 | "libmariadb/ma_init.c", 67 | "libmariadb/ma_password.c", 68 | "libmariadb/ma_ll2str.c", 69 | "libmariadb/ma_sha1.c", 70 | "libmariadb/mariadb_stmt.c", 71 | "libmariadb/ma_loaddata.c", 72 | "libmariadb/ma_stmt_codec.c", 73 | "libmariadb/ma_string.c", 74 | "libmariadb/ma_dtoa.c", 75 | "libmariadb/mariadb_dyncol.c", 76 | "libmariadb/mariadb_async.c", 77 | "libmariadb/ma_context.c", 78 | "libmariadb/ma_io.c", 79 | ], 80 | hdrs = configure_out_hdrs + glob([ 81 | "include/*.h", 82 | "include/mysql/*.h", 83 | "include/mariadb/*.h", 84 | ]), 85 | copts = [ 86 | "-g", 87 | "-DLIBMARIADB", 88 | "-DTHREAD", 89 | "-DHAVE_COMPRESS", 90 | "-DENABLED_LOCAL_INFILE", 91 | "-DLIBICONV_PLUG", 92 | "-DHAVE_OPENSSL", 93 | "-DHAVE_TLS", 94 | ], 95 | includes = [ 96 | "build/include/", 97 | "include/", 98 | "include/mariadb/", 99 | ], 100 | linkopts = [ 101 | "-lpthread", 102 | "-ldl", 103 | "-lm", 104 | ] + select({ 105 | ":darwin": ["-liconv"], 106 | "//conditions:default": [], 107 | }), 108 | visibility = ["//visibility:public"], 109 | deps = [ 110 | "@boringssl//:ssl", 111 | "@zlib", 112 | ], 113 | ) 114 | -------------------------------------------------------------------------------- /ml_metadata/libmysqlclient.WORKSPACE: -------------------------------------------------------------------------------- 1 | workspace(name = "libmysqlclient") 2 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Init module for ML Metadata.""" 15 | from ml_metadata.metadata_store.metadata_store import ( 16 | ListOptions, 17 | MetadataStore, 18 | OrderByField, 19 | downgrade_schema, 20 | ) 21 | 22 | __all__ = [ 23 | "downgrade_schema", 24 | "ListOptions", 25 | "MetadataStore", 26 | "OrderByField", 27 | ] 28 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/conftest.py: -------------------------------------------------------------------------------- 1 | def pytest_addoption(parser): 2 | parser.addoption( 3 | '--use_grpc_backend', action='store_true', default=False, 4 | help='Set this to true to use gRPC instead of sqlLite backend.' 5 | ) 6 | parser.addoption( 7 | '--grpc_host', type=str, default=None, 8 | help="The gRPC host name to use when use_grpc_backed is set to 'True'" 9 | ) 10 | parser.addoption( 11 | '--grpc_port', type=int, default=0, 12 | help="The gRPC port number to use when use_grpc_backed is set to 'True'" 13 | ) 14 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/constants.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_CONSTANTS_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_CONSTANTS_H_ 17 | 18 | #include "absl/strings/string_view.h" 19 | namespace ml_metadata { 20 | 21 | // The in-memory encoding of the NULL value in RecordSet proto returned from 22 | // any MetadataSource. 23 | static constexpr absl::string_view kMetadataSourceNull = "__MLMD_NULL__"; 24 | 25 | // The node type_kind enum values used for internal storage. The enum value 26 | // should not be modified, in order to be backward compatible with stored types. 27 | // LINT.IfChange 28 | enum class TypeKind { EXECUTION_TYPE = 0, ARTIFACT_TYPE = 1, CONTEXT_TYPE = 2 }; 29 | // LINT.ThenChange(../util/metadata_source_query_config.cc) 30 | 31 | // Default maximum number of returned resources for List operation. 32 | constexpr int kDefaultMaxListOperationResultSize = 100; 33 | 34 | constexpr int kPropertyRecordSetSize = 8; 35 | 36 | constexpr int kQueryLineageSubgraphMaxNumHops = 100; 37 | 38 | } // namespace ml_metadata 39 | 40 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_CONSTANTS_H_ 41 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/list_operation_query_helper.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_LIST_OPERATION_QUERY_HELPER_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_LIST_OPERATION_QUERY_HELPER_H_ 17 | 18 | #include "absl/status/status.h" 19 | #include "absl/types/optional.h" 20 | #include "ml_metadata/metadata_store/constants.h" 21 | #include "ml_metadata/metadata_store/types.h" 22 | #include "ml_metadata/proto/metadata_store.pb.h" 23 | 24 | namespace ml_metadata { 25 | 26 | // Utility methods to generate SQL queries for List Operations. 27 | 28 | // Generates the WHERE clause for ListOperation. 29 | // On success `sql_query_clause` is appended with the constructed parameterized 30 | // WHERE clause based on |options| and qualifying column names with 31 | // |table_alias|. 32 | // For the first page of results based on the ordering field this method 33 | // generates a ordering clause based only on that field, e.g., 34 | // `create_time_since_epoch` <= INT_MAX (for desc) and 35 | // `create_time_since_epoch` >= 0 (for asc). 36 | // For following pages method first decodes the next_page_token set in `options` 37 | // and validates the token. 38 | // Based on the ordering field the generated clause is as follows: 39 | // 1. CREATE_TIME 40 | // `create_time_since_epoch` <= |options.field_offset| AND 41 | // `id` < |options.id_offset|` 42 | // 2. LAST_UPDATE_TIME 43 | // `last_update_time_since_epoch` <= |options.field_offset| AND 44 | // `id` NOT IN(options.list_ids)` 45 | // list_ids - Is the list of trailing node ids in the previous page that have 46 | // the same last_update_time_since_epoch value. 47 | // 3. ID 48 | // `id < |options.field_offset|. 49 | // 50 | // Returns INVALID_ARGUMENT error if the `options` or `next_page_token` 51 | // specified is invalid. 52 | absl::Status AppendOrderingThresholdClause( 53 | const ListOperationOptions& options, 54 | std::optional table_alias, 55 | std::string& sql_query_clause); 56 | 57 | // Generates the ORDER BY clause for ListOperation. 58 | // On success `sql_query_clause` is appended with the constructed ORDER BY 59 | // clause based on |options| and qualifying column names with |table_alias|. 60 | // Returns INVALID_ARGUMENT error if the |options| 61 | // specified is invalid. 62 | // For example, given a ListOperationOptions message: 63 | // { 64 | // max_result_size: 1, 65 | // order_by_field: { 66 | // field: CREATE_TIME, 67 | // is_asc: false 68 | // } 69 | // } 70 | // Appends "ORDER BY `create_time_since_epoch` DESC, `id` DESC" at the end of 71 | // `sql_query_clause`. 72 | absl::Status AppendOrderByClause(const ListOperationOptions& options, 73 | std::optional table_alias, 74 | std::string& sql_query_clause); 75 | 76 | // Generates the LIMIT clause for ListOperation. 77 | // On success `sql_query_clause` is appended with the constructed LIMIT clause 78 | // based on |options|. 79 | // If |options| does not specify max_result_size a default value of 20 is used 80 | // and if the max_result_size is greater than 100, the LIMIT clause coerces the 81 | // value to 100. 82 | // For example, given a ListOperationOptions message: 83 | // { 84 | // max_result_size: 1, 85 | // order_by_field: { 86 | // field: CREATE_TIME, 87 | // is_asc: false 88 | // } 89 | // } 90 | // Appends "LIMIT 1" at the end of `sql_query_clause`. 91 | absl::Status AppendLimitClause(const ListOperationOptions& options, 92 | std::string& sql_query_clause); 93 | 94 | // Gets the maximum number of returned resources for List operation. 95 | inline constexpr int GetDefaultMaxListOperationResultSize() { 96 | return kDefaultMaxListOperationResultSize; 97 | } 98 | 99 | } // namespace ml_metadata 100 | 101 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_LIST_OPERATION_QUERY_HELPER_H_ 102 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/list_operation_util.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/list_operation_util.h" 16 | #include 17 | #include 18 | #include 19 | 20 | #include "absl/status/status.h" 21 | #include "absl/strings/escaping.h" 22 | #include "absl/strings/str_cat.h" 23 | #include "absl/strings/string_view.h" 24 | #include "ml_metadata/proto/metadata_store.pb.h" 25 | 26 | namespace ml_metadata { 27 | 28 | // Helper method to set initial field and id offset values based on List 29 | // operation options. 30 | void SetListOperationInitialValues(const ListOperationOptions& options, 31 | int64_t& field_offset, int64_t& id_offset) { 32 | const bool is_asc = options.order_by_field().is_asc(); 33 | field_offset = is_asc ? LLONG_MIN : LLONG_MAX; 34 | id_offset = field_offset; 35 | } 36 | 37 | absl::Status DecodeListOperationNextPageToken( 38 | absl::string_view next_page_token, 39 | ListOperationNextPageToken& list_operation_next_page_token) { 40 | std::string token_str; 41 | if (!absl::WebSafeBase64Unescape(next_page_token, &token_str)) { 42 | return absl::InvalidArgumentError( 43 | "Failed to decode next page token string"); 44 | } 45 | 46 | if (!list_operation_next_page_token.ParseFromString(token_str)) { 47 | return absl::InvalidArgumentError( 48 | "Failed to parse decoded next page token into " 49 | "ListOperationNextPageToken proto message "); 50 | } 51 | 52 | return absl::OkStatus(); 53 | } 54 | 55 | absl::Status ValidateListOperationOptionsAreIdentical( 56 | const ListOperationOptions& previous_options, 57 | const ListOperationOptions& current_options) { 58 | if (previous_options.order_by_field().is_asc() == 59 | current_options.order_by_field().is_asc() && 60 | previous_options.order_by_field().field() == 61 | current_options.order_by_field().field()) { 62 | return absl::OkStatus(); 63 | } 64 | 65 | return absl::InvalidArgumentError(absl::StrCat( 66 | "ListOperationOptions mismatch between calls Initial Options: ", 67 | previous_options.DebugString(), 68 | " Current Options: ", current_options.DebugString())); 69 | } 70 | } // namespace ml_metadata 71 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/list_operation_util.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_LIST_OPERATION_UTIL_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_LIST_OPERATION_UTIL_H_ 17 | 18 | #include "absl/status/status.h" 19 | #include "absl/strings/escaping.h" 20 | #include "absl/types/span.h" 21 | #include "ml_metadata/metadata_store/types.h" 22 | #include "ml_metadata/proto/metadata_store.pb.h" 23 | 24 | namespace ml_metadata { 25 | 26 | // Sets initial field and id offset values based on List operation options. 27 | // TODO(b/187082552): Remove this method. 28 | void SetListOperationInitialValues(const ListOperationOptions& options, 29 | int64_t& field_offset, int64_t& id_offset); 30 | 31 | // Decodes ListOperationNextPageToken encoded in `next_page_token`. 32 | absl::Status DecodeListOperationNextPageToken( 33 | absl::string_view next_page_token, 34 | ListOperationNextPageToken& list_operation_next_page_token); 35 | 36 | // Generates encoded list operation next page token string. 37 | template 38 | absl::Status BuildListOperationNextPageToken( 39 | absl::Span nodes, const ListOperationOptions& options, 40 | std::string* next_page_token) { 41 | const Node& last_node = nodes.back(); 42 | ListOperationNextPageToken list_operation_next_page_token; 43 | switch (options.order_by_field().field()) { 44 | case ListOperationOptions::OrderByField::CREATE_TIME: { 45 | list_operation_next_page_token.set_field_offset( 46 | last_node.create_time_since_epoch()); 47 | list_operation_next_page_token.set_id_offset(last_node.id()); 48 | break; 49 | } 50 | case ListOperationOptions::OrderByField::LAST_UPDATE_TIME: { 51 | list_operation_next_page_token.add_listed_ids(last_node.id()); 52 | list_operation_next_page_token.set_field_offset( 53 | last_node.last_update_time_since_epoch()); 54 | for (int i = nodes.size() - 2; i >= 0; i--) { 55 | if (nodes[i].last_update_time_since_epoch() != 56 | last_node.last_update_time_since_epoch()) { 57 | break; 58 | } 59 | list_operation_next_page_token.add_listed_ids(nodes[i].id()); 60 | } 61 | break; 62 | } 63 | case ListOperationOptions::OrderByField::ID: { 64 | list_operation_next_page_token.set_field_offset(last_node.id()); 65 | list_operation_next_page_token.set_id_offset(last_node.id()); 66 | break; 67 | } 68 | default: 69 | return absl::InvalidArgumentError( 70 | absl::StrCat("Unsupported field: ", 71 | ListOperationOptions::OrderByField::Field_Name( 72 | options.order_by_field().field()), 73 | " specified in ListOperationOptions")); 74 | } 75 | *list_operation_next_page_token.mutable_set_options() = options; 76 | // Clear previous `next_page_token` as it is changing for each request and it 77 | // is a set-only field. If not clean this field, the `next_page_token` in 78 | // the follow-up request encodes itself and the size grows exponentially. 79 | list_operation_next_page_token.mutable_set_options()->clear_next_page_token(); 80 | *next_page_token = absl::WebSafeBase64Escape( 81 | list_operation_next_page_token.SerializeAsString()); 82 | return absl::OkStatus(); 83 | } 84 | 85 | // Ensures that ListOperationOptions have not changed between 86 | // calls. |previous_options| represents options used in the previous call and 87 | // |current_options| represents options used in the current call. 88 | // Validation validates order_by_fields and filter_query in 89 | // ListOperationOptions. 90 | absl::Status ValidateListOperationOptionsAreIdentical( 91 | const ListOperationOptions& previous_options, 92 | const ListOperationOptions& current_options); 93 | } // namespace ml_metadata 94 | 95 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_LIST_OPERATION_UTIL_H_ 96 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/metadata_access_object_factory.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_ACCESS_OBJECT_FACTORY_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_ACCESS_OBJECT_FACTORY_H_ 17 | 18 | #include 19 | #include 20 | #include 21 | 22 | #include "absl/status/status.h" 23 | #include "ml_metadata/metadata_store/metadata_access_object.h" 24 | #include "ml_metadata/metadata_store/metadata_source.h" 25 | #include "ml_metadata/proto/metadata_source.pb.h" 26 | 27 | namespace ml_metadata { 28 | 29 | // Factory method, if the return value is ok, 'result' is populated with an 30 | // object that can be used to access metadata with the given config and 31 | // metadata_source. The caller is responsible to own a MetadataSource, and the 32 | // MetadataAccessObject connects and execute queries with the MetadataSource. 33 | // Returns INVALID_ARGUMENT error, if query_config is not valid. 34 | // Returns detailed INTERNAL error, if the MetadataSource cannot be connected. 35 | absl::Status CreateMetadataAccessObject( 36 | const MetadataSourceQueryConfig& query_config, 37 | MetadataSource* const metadata_source, 38 | std::unique_ptr* result); 39 | 40 | // For multi-tenant applications to have better availability when configured 41 | // with a set of existing backends with different schema versions, an 42 | // `schema_version` can be set to allow the MetadataAccessObject 43 | // works with an existing db having that particular schema version. 44 | absl::Status CreateMetadataAccessObject( 45 | const MetadataSourceQueryConfig& query_config, 46 | MetadataSource* metadata_source, std::optional schema_version, 47 | std::unique_ptr* result); 48 | 49 | } // namespace ml_metadata 50 | 51 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_ACCESS_OBJECT_FACTORY_H_ 52 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/metadata_source.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/metadata_source.h" 16 | 17 | #include "absl/status/status.h" 18 | #include "ml_metadata/util/return_utils.h" 19 | 20 | namespace ml_metadata { 21 | 22 | absl::Status MetadataSource::Connect() { 23 | if (is_connected_) 24 | return absl::FailedPreconditionError( 25 | "The connection has been opened. Close() the current connection before " 26 | "Connect() again."); 27 | MLMD_RETURN_IF_ERROR(ConnectImpl()); 28 | is_connected_ = true; 29 | return absl::OkStatus(); 30 | } 31 | 32 | absl::Status MetadataSource::Close() { 33 | if (!is_connected_) 34 | return absl::FailedPreconditionError( 35 | "No connection is opened when calling Close()."); 36 | MLMD_RETURN_IF_ERROR(CloseImpl()); 37 | is_connected_ = false; 38 | return absl::OkStatus(); 39 | } 40 | 41 | absl::Status MetadataSource::ExecuteQuery(const std::string& query, 42 | RecordSet* results) { 43 | if (!is_connected_) 44 | return absl::FailedPreconditionError("No opened connection for querying."); 45 | if (!transaction_open_) 46 | return absl::FailedPreconditionError("Transaction not open."); 47 | return ExecuteQueryImpl(query, results); 48 | } 49 | 50 | absl::Status MetadataSource::Begin() { 51 | if (!is_connected_) 52 | return absl::FailedPreconditionError("No opened connection for querying."); 53 | if (transaction_open_) 54 | return absl::FailedPreconditionError("Transaction already open."); 55 | MLMD_RETURN_IF_ERROR(BeginImpl()); 56 | transaction_open_ = true; 57 | return absl::OkStatus(); 58 | } 59 | 60 | 61 | absl::Status MetadataSource::Commit() { 62 | if (!is_connected_) 63 | return absl::FailedPreconditionError("No opened connection for querying."); 64 | if (!transaction_open_) 65 | return absl::FailedPreconditionError("Transaction not open."); 66 | MLMD_RETURN_IF_ERROR(CommitImpl()); 67 | transaction_open_ = false; 68 | return absl::OkStatus(); 69 | } 70 | 71 | absl::Status MetadataSource::Rollback() { 72 | if (!is_connected_) 73 | return absl::FailedPreconditionError("No opened connection for querying."); 74 | if (!transaction_open_) 75 | return absl::FailedPreconditionError("Transaction not open."); 76 | MLMD_RETURN_IF_ERROR(RollbackImpl()); 77 | transaction_open_ = false; 78 | return absl::OkStatus(); 79 | } 80 | 81 | } // namespace ml_metadata 82 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/metadata_source_test_suite.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_SOURCE_TEST_SUITE_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_SOURCE_TEST_SUITE_H_ 17 | 18 | #include 19 | #include 20 | 21 | #include 22 | #include 23 | #include "ml_metadata/metadata_store/metadata_source.h" 24 | #include "ml_metadata/proto/metadata_source.pb.h" 25 | 26 | namespace ml_metadata { 27 | namespace testing { 28 | 29 | // An Interface to create, retrieve and clean a MetadataSource. 30 | class MetadataSourceContainer { 31 | public: 32 | MetadataSourceContainer() = default; 33 | virtual ~MetadataSourceContainer() = default; 34 | 35 | // MetadataSource is owned by MetadataSourceContainer. 36 | virtual MetadataSource* GetMetadataSource() = 0; 37 | 38 | // Creates a new table schema for testing. 39 | virtual void InitTestSchema() = 0; 40 | 41 | // Creates table schema and populates rows. 42 | virtual void InitSchemaAndPopulateRows() = 0; 43 | }; 44 | 45 | // Represents the type of the Gunit Test param for the parameterized 46 | // MetadataSourceTestSuite. 47 | // 48 | // Note: Since the Gunit Test param needs to be a copyable type, we use a 49 | // std::function as opposed to directly using 50 | // std::unique_ptr as the Gunit Test param type. 51 | using MetadataSourceContainerFactory = 52 | std::function()>; 53 | 54 | class MetadataSourceTestSuite 55 | : public ::testing::TestWithParam { 56 | protected: 57 | void SetUp() override { 58 | metadata_source_container_ = GetParam()(); 59 | metadata_source_ = metadata_source_container_->GetMetadataSource(); 60 | } 61 | 62 | void TearDown() override { 63 | metadata_source_ = nullptr; 64 | metadata_source_container_ = nullptr; 65 | } 66 | 67 | std::unique_ptr metadata_source_container_; 68 | // metadata_source_ is unowned. 69 | MetadataSource* metadata_source_; 70 | }; 71 | 72 | } // namespace testing 73 | } // namespace ml_metadata 74 | 75 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_SOURCE_TEST_SUITE_H_ 76 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/metadata_store_factory.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_STORE_FACTORY_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_STORE_FACTORY_H_ 17 | 18 | #include 19 | 20 | #include "absl/status/status.h" 21 | #include "ml_metadata/metadata_store/metadata_store.h" 22 | #include "ml_metadata/proto/metadata_store.pb.h" 23 | 24 | namespace ml_metadata { 25 | // Creates a MetadataStore. 26 | // If the method returns OK, the method MUST set result to contain 27 | // a non-null pointer. 28 | absl::Status CreateMetadataStore(const ConnectionConfig& config, 29 | std::unique_ptr* result); 30 | 31 | absl::Status CreateMetadataStore(const ConnectionConfig& config, 32 | const MigrationOptions& options, 33 | std::unique_ptr* result); 34 | 35 | } // namespace ml_metadata 36 | 37 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_STORE_FACTORY_H_ 38 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/metadata_store_factory_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/metadata_store_factory.h" 16 | 17 | #include 18 | 19 | #include 20 | #include 21 | #include 22 | 23 | #include "ml_metadata/metadata_store/metadata_access_object.h" 24 | #include "ml_metadata/metadata_store/metadata_store.h" 25 | #include "ml_metadata/metadata_store/test_util.h" 26 | #include "ml_metadata/proto/metadata_source.pb.h" 27 | #include "ml_metadata/proto/metadata_store.pb.h" 28 | 29 | namespace ml_metadata { 30 | namespace testing { 31 | using testing::ParseTextProtoOrDie; 32 | 33 | // Utility method to test CreateMetadataStore. 34 | // Given a ConnectionConfig, it creates the store, puts and gets an 35 | // artifact type. 36 | void TestPutAndGetArtifactType(const ConnectionConfig& connection_config) { 37 | std::unique_ptr store; 38 | const MigrationOptions opts; 39 | ASSERT_EQ(absl::OkStatus(), 40 | CreateMetadataStore(connection_config, opts, &store)); 41 | 42 | PutArtifactTypeRequest put_request = 43 | ParseTextProtoOrDie( 44 | R"( 45 | all_fields_match: true 46 | artifact_type: { 47 | name: 'test_type2' 48 | properties { key: 'property_1' value: STRING } 49 | } 50 | )"); 51 | PutArtifactTypeResponse put_response; 52 | ASSERT_EQ(absl::OkStatus(), 53 | store->PutArtifactType(put_request, &put_response)); 54 | ASSERT_TRUE(put_response.has_type_id()); 55 | GetArtifactTypeRequest get_request = 56 | ParseTextProtoOrDie( 57 | R"( 58 | type_name: 'test_type2' 59 | )"); 60 | GetArtifactTypeResponse get_response; 61 | ASSERT_EQ(absl::OkStatus(), 62 | store->GetArtifactType(get_request, &get_response)); 63 | EXPECT_THAT(get_response.artifact_type(), 64 | testing::EqualsProto(put_request.artifact_type(), 65 | /*ignore_fields=*/{"id"})) 66 | << "The type should be the same as the one given."; 67 | } 68 | 69 | 70 | TEST(MetadataStoreFactoryTest, CreateSQLiteMetadataStore) { 71 | ConnectionConfig connection_config; 72 | connection_config.mutable_sqlite(); 73 | TestPutAndGetArtifactType(connection_config); 74 | } 75 | 76 | } // namespace testing 77 | } // namespace ml_metadata 78 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/metadata_store_test_suite.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_STORE_TEST_SUITE_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_STORE_TEST_SUITE_H_ 17 | 18 | #include 19 | #include 20 | 21 | #include 22 | #include 23 | #include "ml_metadata/metadata_store/metadata_store.h" 24 | 25 | namespace ml_metadata { 26 | namespace testing { 27 | 28 | // An Interface to create, retrieve and clean a MetadataStore object. 29 | class MetadataStoreContainer { 30 | public: 31 | MetadataStoreContainer() = default; 32 | virtual ~MetadataStoreContainer() = default; 33 | 34 | // MetadataStore is owned by MetadataStoreContainer. 35 | virtual MetadataStore* GetMetadataStore() = 0; 36 | }; 37 | 38 | // Represents the type of the Gunit Test param for the parameterized 39 | // MetadataStoreTestSuite. 40 | // 41 | // Note: Since the Gunit Test param needs to be a copyable type, we use a 42 | // std::function as opposed to directly using 43 | // std::unique_ptr as the Gunit Test param type. 44 | using MetadataStoreContainerFactory = 45 | std::function()>; 46 | 47 | class MetadataStoreTestSuite 48 | : public ::testing::TestWithParam { 49 | protected: 50 | void SetUp() override { 51 | metadata_store_container_ = GetParam()(); 52 | metadata_store_ = metadata_store_container_->GetMetadataStore(); 53 | } 54 | 55 | void TearDown() override { 56 | metadata_store_ = nullptr; 57 | metadata_store_container_ = nullptr; 58 | } 59 | 60 | std::unique_ptr metadata_store_container_; 61 | MetadataStore* metadata_store_; 62 | }; 63 | 64 | } // namespace testing 65 | } // namespace ml_metadata 66 | 67 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_METADATA_STORE_TEST_SUITE_H_ 68 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/mlmd_types_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Tests for mlmd_types.""" 15 | from absl.testing import absltest 16 | 17 | from ml_metadata.metadata_store import mlmd_types 18 | from ml_metadata.proto import metadata_store_pb2 19 | 20 | 21 | def _get_artifact_type_name( 22 | system_type: metadata_store_pb2.ArtifactType.SystemDefinedBaseType): 23 | extensions = metadata_store_pb2.ArtifactType.SystemDefinedBaseType.DESCRIPTOR.values_by_number[ 24 | system_type].GetOptions().Extensions 25 | return extensions[metadata_store_pb2.system_type_extension].type_name 26 | 27 | 28 | def _get_execution_type_name( 29 | system_type: metadata_store_pb2.ExecutionType.SystemDefinedBaseType): 30 | extensions = metadata_store_pb2.ExecutionType.SystemDefinedBaseType.DESCRIPTOR.values_by_number[ 31 | system_type].GetOptions().Extensions 32 | return extensions[metadata_store_pb2.system_type_extension].type_name 33 | 34 | 35 | class MlmdTypesTest(absltest.TestCase): 36 | 37 | def testSystemArtifactTypes(self): 38 | self.assertEqual( 39 | _get_artifact_type_name(metadata_store_pb2.ArtifactType.DATASET), 40 | mlmd_types.Dataset().name) 41 | self.assertEqual( 42 | _get_artifact_type_name(metadata_store_pb2.ArtifactType.MODEL), 43 | mlmd_types.Model().name) 44 | self.assertEqual( 45 | _get_artifact_type_name(metadata_store_pb2.ArtifactType.METRICS), 46 | mlmd_types.Metrics().name) 47 | self.assertEqual( 48 | _get_artifact_type_name(metadata_store_pb2.ArtifactType.STATISTICS), 49 | mlmd_types.Statistics().name) 50 | 51 | def testSystemExecutionTypes(self): 52 | self.assertEqual( 53 | _get_execution_type_name(metadata_store_pb2.ExecutionType.TRAIN), 54 | mlmd_types.Train().name) 55 | self.assertEqual( 56 | _get_execution_type_name(metadata_store_pb2.ExecutionType.TRANSFORM), 57 | mlmd_types.Transform().name) 58 | self.assertEqual( 59 | _get_execution_type_name(metadata_store_pb2.ExecutionType.PROCESS), 60 | mlmd_types.Process().name) 61 | self.assertEqual( 62 | _get_execution_type_name(metadata_store_pb2.ExecutionType.EVALUATE), 63 | mlmd_types.Evaluate().name) 64 | self.assertEqual( 65 | _get_execution_type_name(metadata_store_pb2.ExecutionType.DEPLOY), 66 | mlmd_types.Deploy().name) 67 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/mysql_metadata_access_object_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | // Test suite for a MySqlMetadataSource based MetadataAccessObject. 16 | 17 | #include 18 | #include 19 | #include 20 | 21 | #include 22 | #include 23 | #include 24 | #include "absl/memory/memory.h" 25 | #include "absl/status/status.h" 26 | #include "absl/types/optional.h" 27 | #include "ml_metadata/metadata_store/metadata_access_object_factory.h" 28 | #include "ml_metadata/metadata_store/metadata_access_object_test.h" 29 | #include "ml_metadata/metadata_store/metadata_source.h" 30 | #include "ml_metadata/metadata_store/mysql_metadata_source.h" 31 | #include "ml_metadata/metadata_store/test_mysql_metadata_source_initializer.h" 32 | #include "ml_metadata/proto/metadata_source.pb.h" 33 | #include "ml_metadata/util/metadata_source_query_config.h" 34 | 35 | namespace ml_metadata { 36 | namespace testing { 37 | 38 | namespace { 39 | 40 | // MySqlMetadataAccessObjectContainer implements MetadataAccessObjectContainer 41 | // to generate and retrieve a MetadataAccessObject based off a 42 | // MySqlMetadataSource. 43 | class MySqlMetadataAccessObjectContainer 44 | : public QueryConfigMetadataAccessObjectContainer { 45 | public: 46 | MySqlMetadataAccessObjectContainer( 47 | std::optional earlier_schema_version = absl::nullopt) 48 | : QueryConfigMetadataAccessObjectContainer( 49 | util::GetMySqlMetadataSourceQueryConfig(), earlier_schema_version) { 50 | metadata_source_initializer_ = GetTestMySqlMetadataSourceInitializer(); 51 | metadata_source_ = metadata_source_initializer_->Init( 52 | TestMySqlMetadataSourceInitializer::ConnectionType::kTcp); 53 | CHECK_EQ(absl::OkStatus(), 54 | CreateMetadataAccessObject( 55 | util::GetMySqlMetadataSourceQueryConfig(), metadata_source_, 56 | earlier_schema_version, &metadata_access_object_)); 57 | } 58 | 59 | ~MySqlMetadataAccessObjectContainer() override { 60 | metadata_source_initializer_->Cleanup(); 61 | }; 62 | 63 | MetadataSource* GetMetadataSource() override { return metadata_source_; } 64 | MetadataAccessObject* GetMetadataAccessObject() override { 65 | return metadata_access_object_.get(); 66 | } 67 | 68 | std::string GetTableNumQuery() final { 69 | return "select count(*) from `information_schema`.`tables` where " 70 | "`table_schema`=(SELECT DATABASE())"; 71 | } 72 | 73 | std::string GetIndexNumQuery() final { 74 | return "select count(*) from `information_schema`.`statistics` where " 75 | "`table_schema`=(SELECT DATABASE())"; 76 | } 77 | 78 | private: 79 | // An unowned TestMySqlMetadataSourceInitializer from a call to 80 | // GetTestMySqlMetadataSourceInitializer(). 81 | std::unique_ptr 82 | metadata_source_initializer_; 83 | // An unowned MySqlMetadataSource from a call to 84 | // metadata_source_initializer->Init(). 85 | MySqlMetadataSource* metadata_source_; 86 | std::unique_ptr metadata_access_object_; 87 | }; 88 | 89 | } // namespace 90 | 91 | INSTANTIATE_TEST_SUITE_P( 92 | MySqlMetadataAccessObjectTest, MetadataAccessObjectTest, 93 | ::testing::Values( 94 | []() { return std::make_unique(); }, 95 | // TODO(b/257334039) Cleanup after V10+ migration 96 | []() { 97 | return std::make_unique( 98 | /*earlier_schema_version=*/9); 99 | }, 100 | []() { 101 | return std::make_unique( 102 | /*earlier_schema_version=*/8); 103 | }, 104 | []() { 105 | return std::make_unique( 106 | /*earlier_schema_version=*/7); 107 | })); 108 | 109 | } // namespace testing 110 | } // namespace ml_metadata 111 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/mysql_query_config_executor_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | // Test suite for a mysql query config-based QueryExecutor. 16 | #include 17 | 18 | #include 19 | #include 20 | #include 21 | #include "absl/memory/memory.h" 22 | #include "ml_metadata/metadata_store/query_config_executor.h" 23 | #include "ml_metadata/metadata_store/query_executor.h" 24 | #include "ml_metadata/metadata_store/query_executor_test.h" 25 | #include "ml_metadata/metadata_store/test_mysql_metadata_source_initializer.h" 26 | #include "ml_metadata/proto/metadata_source.pb.h" 27 | #include "ml_metadata/util/metadata_source_query_config.h" 28 | 29 | namespace ml_metadata { 30 | namespace testing { 31 | 32 | namespace { 33 | // MySqlQueryConfigExecutorContainer implements 34 | // QueryConfigExecutorContainer to generate and retrieve a 35 | // QueryExecutor based on a MySqlMetadataSource. 36 | class MySqlQueryConfigExecutorContainer : public QueryConfigExecutorContainer { 37 | public: 38 | MySqlQueryConfigExecutorContainer() 39 | : QueryConfigExecutorContainer( 40 | util::GetMySqlMetadataSourceQueryConfig()) { 41 | metadata_source_initializer_ = GetTestMySqlMetadataSourceInitializer(); 42 | metadata_source_ = metadata_source_initializer_->Init( 43 | TestMySqlMetadataSourceInitializer::ConnectionType::kTcp); 44 | if (!metadata_source_->is_connected()) 45 | CHECK_EQ(absl::OkStatus(), metadata_source_->Connect()); 46 | query_executor_ = absl::WrapUnique(new QueryConfigExecutor( 47 | util::GetMySqlMetadataSourceQueryConfig(), metadata_source_)); 48 | } 49 | 50 | ~MySqlQueryConfigExecutorContainer() override = default; 51 | 52 | MetadataSource* GetMetadataSource() override { 53 | return metadata_source_; 54 | } 55 | QueryExecutor* GetQueryExecutor() override { return query_executor_.get(); } 56 | 57 | private: 58 | // An unowned TestMySqlMetadataSourceInitializer from a call to 59 | // GetTestMySqlMetadataSourceInitializer(). 60 | std::unique_ptr 61 | metadata_source_initializer_; 62 | // An unowned MySqlMetadataSource from a call to 63 | // metadata_source_initializer->Init(). 64 | MySqlMetadataSource* metadata_source_; 65 | std::unique_ptr query_executor_; 66 | }; 67 | 68 | } // namespace 69 | 70 | INSTANTIATE_TEST_SUITE_P( 71 | MySqlQueryConfigExecutorContainer, QueryExecutorTest, 72 | ::testing::Values([]() { 73 | return std::make_unique(); 74 | })); 75 | 76 | } // namespace testing 77 | } // namespace ml_metadata 78 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/postgresql_metadata_source_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2023 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | 16 | #include "ml_metadata/metadata_store/postgresql_metadata_source.h" 17 | 18 | #include "gflags/gflags.h" 19 | #include 20 | #include 21 | #include "ml_metadata/metadata_store/metadata_source_test_suite.h" 22 | #include "ml_metadata/metadata_store/test_postgresql_metadata_source_initializer.h" 23 | 24 | namespace ml_metadata { 25 | namespace testing { 26 | namespace { 27 | 28 | // This container calls initializer to set up PostgreSQL database. After DB 29 | // becomes available, this container will get metadata source from initializer. 30 | // Then schema for testing can be initialized by this container. 31 | class PostgreSQLMetadataSourceContainer : public MetadataSourceContainer { 32 | public: 33 | PostgreSQLMetadataSourceContainer() : MetadataSourceContainer() { 34 | metadata_source_initializer_ = GetTestPostgreSQLMetadataSourceInitializer(); 35 | metadata_source_ = metadata_source_initializer_->Init(); 36 | } 37 | 38 | ~PostgreSQLMetadataSourceContainer() override { 39 | metadata_source_initializer_->Cleanup(); 40 | } 41 | 42 | MetadataSource* GetMetadataSource() override { return metadata_source_; } 43 | 44 | // Creates a new table t1(c1 INT, c2 VARCHAR(255)). 45 | void InitTestSchema() override { 46 | CHECK_EQ(metadata_source_->Connect(), absl::OkStatus()); 47 | CHECK_EQ(metadata_source_->Begin(), absl::OkStatus()); 48 | CHECK_EQ(metadata_source_->ExecuteQuery( 49 | "CREATE TABLE t1 (c1 INT, c2 VARCHAR(255));", nullptr), 50 | absl::OkStatus()); 51 | CHECK_EQ(metadata_source_->Commit(), absl::OkStatus()); 52 | } 53 | 54 | // Creates table t1(c1 INT, c2 VARCHAR(255)) and 55 | // adds 3 rows to this table: (1,'v1'), (2,'v2'), (3, 'v3'). 56 | void InitSchemaAndPopulateRows() override { 57 | InitTestSchema(); 58 | CHECK_EQ(metadata_source_->Begin(), absl::OkStatus()); 59 | CHECK_EQ(metadata_source_->ExecuteQuery("INSERT INTO t1 VALUES (1, 'v1')", 60 | nullptr), 61 | absl::OkStatus()); 62 | CHECK_EQ(metadata_source_->ExecuteQuery("INSERT INTO t1 VALUES (2, 'v2')", 63 | nullptr), 64 | absl::OkStatus()); 65 | CHECK_EQ(metadata_source_->ExecuteQuery("INSERT INTO t1 VALUES (3, 'v3')", 66 | nullptr), 67 | absl::OkStatus()); 68 | CHECK_EQ(metadata_source_->Commit(), absl::OkStatus()); 69 | } 70 | 71 | private: 72 | // A TestPostgreSQLMetadataSourceInitializer from a call to 73 | // GetTestPostgreSQLMetadataSourceInitializer(). 74 | std::unique_ptr 75 | metadata_source_initializer_; 76 | 77 | PostgreSQLMetadataSource* metadata_source_; 78 | }; 79 | 80 | } // namespace 81 | 82 | INSTANTIATE_TEST_SUITE_P( 83 | PostgreSQLMetadataSourceTest, MetadataSourceTestSuite, 84 | ::testing::Values([]() { 85 | return std::make_unique(); 86 | })); 87 | 88 | } // namespace testing 89 | } // namespace ml_metadata 90 | 91 | int main(int argc, char* argv[]) { 92 | ::testing::InitGoogleTest(&argc, argv); 93 | gflags::ParseCommandLineFlags(&argc, &argv, true); 94 | return RUN_ALL_TESTS(); 95 | } 96 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/postgresql_query_executor_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2023 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | // Test suite for a PostgreSQL query config-based QueryExecutor. 16 | #include "ml_metadata/metadata_store/postgresql_query_executor.h" 17 | 18 | #include 19 | 20 | #include 21 | #include "absl/memory/memory.h" 22 | #include "ml_metadata/metadata_store/query_executor.h" 23 | #include "ml_metadata/metadata_store/query_executor_test.h" 24 | #include "ml_metadata/metadata_store/test_postgresql_metadata_source_initializer.h" 25 | #include "ml_metadata/util/metadata_source_query_config.h" 26 | 27 | namespace ml_metadata { 28 | namespace testing { 29 | 30 | namespace { 31 | // PostgreSQLQueryExecutorContainer implements 32 | // QueryConfigExecutorContainer to generate and retrieve a 33 | // QueryExecutor based on a PostgreSQLMetadataSource. 34 | class PostgreSQLQueryExecutorContainer : public QueryConfigExecutorContainer { 35 | public: 36 | PostgreSQLQueryExecutorContainer() 37 | : QueryConfigExecutorContainer( 38 | util::GetPostgreSQLMetadataSourceQueryConfig()) { 39 | metadata_source_initializer_ = GetTestPostgreSQLMetadataSourceInitializer(); 40 | metadata_source_ = metadata_source_initializer_->Init(); 41 | if (!metadata_source_->is_connected()) 42 | CHECK_EQ(metadata_source_->Connect(), absl::OkStatus()); 43 | query_executor_ = absl::WrapUnique(new PostgreSQLQueryExecutor( 44 | util::GetPostgreSQLMetadataSourceQueryConfig(), metadata_source_)); 45 | } 46 | 47 | ~PostgreSQLQueryExecutorContainer() override = default; 48 | 49 | MetadataSource* GetMetadataSource() override { return metadata_source_; } 50 | QueryExecutor* GetQueryExecutor() override { return query_executor_.get(); } 51 | 52 | private: 53 | // An unowned TestPostgreSQLMetadataSourceInitializer from a call to 54 | // GetTestPostgreSQLMetadataSourceInitializer(). 55 | std::unique_ptr 56 | metadata_source_initializer_; 57 | // An unowned PostgreSQLMetadataSource from a call to 58 | // metadata_source_initializer->Init(). 59 | PostgreSQLMetadataSource* metadata_source_; 60 | std::unique_ptr query_executor_; 61 | }; 62 | 63 | } // namespace 64 | 65 | INSTANTIATE_TEST_SUITE_P( 66 | PostgreSQLQueryExecutorContainer, QueryExecutorTest, 67 | ::testing::Values([]() { 68 | return std::make_unique(); 69 | })); 70 | 71 | } // namespace testing 72 | } // namespace ml_metadata 73 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/pywrap/BUILD: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # Placeholder: load py_library 16 | load( 17 | "//ml_metadata:ml_metadata.bzl", 18 | "ml_metadata_pybind_extension", 19 | ) 20 | 21 | package(default_visibility = [ 22 | "//ml_metadata:__subpackages__", 23 | ]) 24 | 25 | licenses(["notice"]) 26 | 27 | ml_metadata_pybind_extension( 28 | name = "metadata_store_extension", 29 | srcs = [ 30 | "metadata_store_extension.cc", 31 | ], 32 | module_name = "metadata_store_extension", 33 | deps = [ 34 | "@com_google_absl//absl/status", 35 | "//ml_metadata/metadata_store:metadata_store_factory", 36 | "//ml_metadata/metadata_store:simple_types_util", 37 | "//ml_metadata/proto:metadata_store_proto", 38 | "//ml_metadata/simple_types/proto:simple_types_proto", 39 | "@pybind11", 40 | ], 41 | ) 42 | 43 | py_library( 44 | name = "metadata_store_extension_lib", 45 | srcs = ["__init__.py"], 46 | deps = [ 47 | ":metadata_store_extension", 48 | ], 49 | ) 50 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/pywrap/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/query_executor.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/query_executor.h" 16 | 17 | #include 18 | 19 | #include 20 | #include "absl/status/status.h" 21 | #include "absl/strings/str_cat.h" 22 | #include "ml_metadata/util/return_utils.h" 23 | 24 | namespace ml_metadata { 25 | 26 | // The kSupportedEarlierQueryVersion can be the same with the current library 27 | // schema version or a schema version before hand. The latter is for 28 | // supporting migration for MLMD online services. 29 | static constexpr int64_t kSupportedEarlierQueryVersion = 7; 30 | 31 | QueryExecutor::QueryExecutor(std::optional query_schema_version) 32 | : query_schema_version_(query_schema_version) { 33 | if (query_schema_version_) { 34 | CHECK_GE(*query_schema_version_, kSupportedEarlierQueryVersion) 35 | << "The query config executor does not support other earlier query " 36 | "version other than " 37 | << kSupportedEarlierQueryVersion 38 | << "; query_schema_version: " << *query_schema_version_; 39 | } 40 | } 41 | 42 | absl::Status QueryExecutor::VerifyCurrentQueryVersionIsAtLeast( 43 | int64_t min_schema_version) const { 44 | return query_schema_version_ && *query_schema_version_ < min_schema_version 45 | ? absl::FailedPreconditionError(absl::StrCat( 46 | "The query executor method requires query_schema_version " 47 | ">= ", 48 | min_schema_version, 49 | "; current query version: ", *query_schema_version_)) 50 | : absl::OkStatus(); 51 | } 52 | 53 | absl::Status QueryExecutor::CheckSchemaVersionAlignsWithQueryVersion() { 54 | if (!query_schema_version_) { 55 | return absl::InvalidArgumentError( 56 | "When query_schema_version_ is set, the method checks the given db has " 57 | "already initialized with the schema version."); 58 | } 59 | int64_t db_version = -1; 60 | const absl::Status existing_schema_version_status = 61 | GetSchemaVersion(&db_version); 62 | if (absl::IsNotFound(existing_schema_version_status)) { 63 | return absl::FailedPreconditionError( 64 | "When using the query executor with query_version other than head, " 65 | "the db should already exists. For empty db, init it using head " 66 | "version with the default constructor"); 67 | } 68 | MLMD_RETURN_IF_ERROR(existing_schema_version_status); 69 | if (db_version != *query_schema_version_) { 70 | return absl::FailedPreconditionError(absl::StrCat( 71 | "The query executor is configured with a query_version: ", 72 | *query_schema_version_, 73 | ", however the underlying db is at schema_version: ", db_version)); 74 | } 75 | return absl::OkStatus(); 76 | } 77 | 78 | bool QueryExecutor::IsQuerySchemaVersionEquals(int64_t schema_version) const { 79 | return query_schema_version_ && *query_schema_version_ == schema_version; 80 | } 81 | 82 | } // namespace ml_metadata 83 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/query_executor_test.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_QUERY_EXECUTOR_TEST_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_QUERY_EXECUTOR_TEST_H_ 17 | 18 | #include 19 | 20 | #include 21 | #include 22 | #include 23 | #include "absl/status/status.h" 24 | #include "ml_metadata/metadata_store/metadata_source.h" 25 | #include "ml_metadata/metadata_store/query_executor.h" 26 | #include "ml_metadata/util/return_utils.h" 27 | 28 | namespace ml_metadata { 29 | namespace testing { 30 | 31 | class QueryExecutorContainer { 32 | public: 33 | virtual ~QueryExecutorContainer() = default; 34 | 35 | virtual MetadataSource* GetMetadataSource() = 0; 36 | 37 | virtual QueryExecutor* GetQueryExecutor() = 0; 38 | 39 | virtual absl::Status Init() { 40 | MLMD_RETURN_IF_ERROR(GetQueryExecutor()->InitMetadataSource()); 41 | return absl::OkStatus(); 42 | } 43 | 44 | // Adds a commit point in the tests. 45 | // Default to be a no-op for SQLite, MySQL. 46 | virtual absl::Status AddCommitPoint() { return absl::OkStatus(); } 47 | }; 48 | 49 | class QueryConfigExecutorContainer : public QueryExecutorContainer { 50 | public: 51 | QueryConfigExecutorContainer(const MetadataSourceQueryConfig& config) 52 | : config_(config) {} 53 | 54 | virtual ~QueryConfigExecutorContainer() = default; 55 | 56 | virtual QueryExecutor* GetQueryExecutor() = 0; 57 | 58 | private: 59 | MetadataSourceQueryConfig config_; 60 | }; 61 | 62 | // Represents the type of the Gunit Test param for the parameterized 63 | // QueryExecutorTest. 64 | using QueryExecutorContainerFactory = 65 | std::function()>; 66 | 67 | // A parameterized abstract test fixture to run tests for 68 | // QueryExecutor created with different MetadataSource types. 69 | // See query_executor_test.cc for list of test cases using this 70 | // fixture. 71 | class QueryExecutorTest 72 | : public ::testing::TestWithParam { 73 | protected: 74 | void SetUp() override { 75 | query_executor_container_ = GetParam()(); 76 | metadata_source_ = query_executor_container_->GetMetadataSource(); 77 | query_executor_ = query_executor_container_->GetQueryExecutor(); 78 | CHECK_EQ(absl::OkStatus(), metadata_source_->Begin()); 79 | } 80 | 81 | void TearDown() override { 82 | CHECK_EQ(absl::OkStatus(), metadata_source_->Commit()); 83 | metadata_source_ = nullptr; 84 | query_executor_ = nullptr; 85 | query_executor_container_ = nullptr; 86 | } 87 | 88 | absl::Status Init() { return query_executor_container_->Init(); } 89 | 90 | // Uses to a add commit point if needed in the tests. 91 | // Default to be a no-op for SQLite, MySQL. 92 | absl::Status AddCommitPointIfNeeded() { 93 | return query_executor_container_->AddCommitPoint(); 94 | } 95 | 96 | std::unique_ptr query_executor_container_; 97 | 98 | // metadata_source_ and query_executor_ are unowned. 99 | MetadataSource* metadata_source_; 100 | QueryExecutor* query_executor_; 101 | }; 102 | 103 | } // namespace testing 104 | } // namespace ml_metadata 105 | 106 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_QUERY_EXECUTOR_TEST_H_ 107 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/simple_types_util.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/simple_types_util.h" 16 | 17 | #include "google/protobuf/text_format.h" 18 | #include "absl/container/flat_hash_map.h" 19 | #include "absl/status/status.h" 20 | #include "absl/strings/match.h" 21 | #include "ml_metadata/proto/metadata_store.pb.h" 22 | #include "ml_metadata/simple_types/proto/simple_types.pb.h" 23 | #include "ml_metadata/simple_types/simple_types_constants.h" 24 | 25 | namespace ml_metadata { 26 | 27 | absl::Status LoadSimpleTypes(SimpleTypes& simple_types) { 28 | if (!google::protobuf::TextFormat::ParseFromString(std::string(kSimpleTypes), 29 | &simple_types)) { 30 | return absl::InvalidArgumentError( 31 | "Failed to parse simple types from string"); 32 | } 33 | return absl::OkStatus(); 34 | } 35 | 36 | bool IsUnsetBaseType(const SystemTypeExtension& extension) { 37 | return absl::StartsWith(extension.type_name(), "unset"); 38 | } 39 | 40 | absl::Status GetSystemTypeEnum(const SystemTypeExtension& extension, 41 | ArtifactType::SystemDefinedBaseType& type_enum) { 42 | static const auto& type_name_to_enum = 43 | *new absl::flat_hash_map( 45 | {{"unset_artifact_type", ArtifactType::UNSET}, 46 | {"mlmd.Dataset", ArtifactType::DATASET}, 47 | {"mlmd.Model", ArtifactType::MODEL}, 48 | {"mlmd.Metrics", ArtifactType::METRICS}, 49 | {"mlmd.Statistics", ArtifactType::STATISTICS}}); 50 | 51 | if (type_name_to_enum.contains(extension.type_name())) { 52 | type_enum = type_name_to_enum.at(extension.type_name()); 53 | return absl::OkStatus(); 54 | } 55 | return absl::InvalidArgumentError( 56 | absl::StrCat("invalid system type name: ", extension.type_name())); 57 | } 58 | 59 | absl::Status GetSystemTypeEnum( 60 | const SystemTypeExtension& extension, 61 | ExecutionType::SystemDefinedBaseType& type_enum) { 62 | static const auto& type_name_to_enum = 63 | *new absl::flat_hash_map( 65 | {{"unset_execution_type", ExecutionType::UNSET}, 66 | {"mlmd.Train", ExecutionType::TRAIN}, 67 | {"mlmd.Transform", ExecutionType::TRANSFORM}, 68 | {"mlmd.Process", ExecutionType::PROCESS}, 69 | {"mlmd.Evaluate", ExecutionType::EVALUATE}, 70 | {"mlmd.Deploy", ExecutionType::DEPLOY}}); 71 | 72 | if (type_name_to_enum.contains(extension.type_name())) { 73 | type_enum = type_name_to_enum.at(extension.type_name()); 74 | return absl::OkStatus(); 75 | } 76 | return absl::InvalidArgumentError( 77 | absl::StrCat("invalid system type name: ", extension.type_name())); 78 | } 79 | 80 | } // namespace ml_metadata 81 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/simple_types_util.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_SIMPLE_TYPES_UTIL_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_SIMPLE_TYPES_UTIL_H_ 17 | 18 | #include "google/protobuf/descriptor.pb.h" 19 | #include "absl/status/status.h" 20 | #include "ml_metadata/proto/metadata_store.pb.h" 21 | #include "ml_metadata/simple_types/proto/simple_types.pb.h" 22 | 23 | namespace ml_metadata { 24 | 25 | // A util to create a SimpleTypes proto from a constant string at runtime. 26 | absl::Status LoadSimpleTypes(SimpleTypes& simple_types); 27 | 28 | // Returns the SystemTypeExtension of the enum field base_type in input 'type'. 29 | template 30 | absl::Status GetSystemTypeExtension(const T& type, 31 | SystemTypeExtension& output_extension) { 32 | if (!type.has_base_type()) { 33 | return absl::NotFoundError("base_type field is null"); 34 | } 35 | const std::string enum_field_name = "base_type"; 36 | auto* message_descriptor = type.GetDescriptor(); 37 | auto* field_descriptor = message_descriptor->FindFieldByName(enum_field_name); 38 | if (!field_descriptor) { 39 | return absl::InvalidArgumentError( 40 | absl::StrCat("Field ", enum_field_name, " is missing in ", 41 | message_descriptor->full_name())); 42 | } 43 | const int enum_value = 44 | type.GetReflection()->GetEnumValue(type, field_descriptor); 45 | 46 | auto* enum_descriptor = 47 | field_descriptor->enum_type()->FindValueByNumber(enum_value); 48 | if (!enum_descriptor) { 49 | return absl::InvalidArgumentError( 50 | absl::StrCat("Enum value of ", enum_field_name, " is: ", enum_value, 51 | ". Failed to get its enum descriptor")); 52 | } 53 | 54 | if (!enum_descriptor->options().HasExtension( 55 | ml_metadata::system_type_extension)) { 56 | return absl::InvalidArgumentError(absl::StrCat( 57 | enum_field_name, " does not have extension to enum value options")); 58 | } 59 | output_extension = enum_descriptor->options().GetExtension( 60 | ml_metadata::system_type_extension); 61 | return absl::OkStatus(); 62 | } 63 | 64 | // Asserts whether the symtem type is UNSET based on 'extension'. 65 | bool IsUnsetBaseType(const SystemTypeExtension& extension); 66 | 67 | // Gets SystemDefinedArtifactType enum based on artifact type 'extension'. 68 | absl::Status GetSystemTypeEnum(const SystemTypeExtension& extension, 69 | ArtifactType::SystemDefinedBaseType& type_enum); 70 | 71 | // Gets SystemDefinedExecutionType enum based on execution type 'extension'. 72 | absl::Status GetSystemTypeEnum(const SystemTypeExtension& extension, 73 | ExecutionType::SystemDefinedBaseType& type_enum); 74 | 75 | } // namespace ml_metadata 76 | 77 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_SIMPLE_TYPES_UTIL_H_ 78 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/simple_types_util_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/simple_types_util.h" 16 | 17 | #include 18 | #include 19 | #include "absl/status/status.h" 20 | #include "ml_metadata/proto/metadata_store.pb.h" 21 | 22 | namespace ml_metadata { 23 | namespace testing { 24 | namespace { 25 | 26 | TEST(SimpleTypesUtilTest, GetArtifactSystemTypeExtensionAndEnum) { 27 | ArtifactType type; 28 | std::vector want_enums = { 29 | ArtifactType::DATASET, ArtifactType::MODEL, ArtifactType::METRICS, 30 | ArtifactType::STATISTICS}; 31 | for (const auto& want_enum : want_enums) { 32 | type.set_base_type(want_enum); 33 | SystemTypeExtension extension; 34 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeExtension(type, extension)); 35 | ArtifactType::SystemDefinedBaseType got_enum; 36 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeEnum(extension, got_enum)); 37 | EXPECT_EQ(want_enum, got_enum); 38 | } 39 | } 40 | 41 | TEST(SimpleTypesUtilTest, GetExecutionSystemTypeExtensionAndEnum) { 42 | ExecutionType type; 43 | std::vector want_enums = { 44 | ExecutionType::TRAIN, ExecutionType::EVALUATE, ExecutionType::DEPLOY, 45 | ExecutionType::TRANSFORM, ExecutionType::PROCESS}; 46 | for (const auto& want_enum : want_enums) { 47 | type.set_base_type(want_enum); 48 | SystemTypeExtension extension; 49 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeExtension(type, extension)); 50 | ExecutionType::SystemDefinedBaseType got_enum; 51 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeEnum(extension, got_enum)); 52 | EXPECT_EQ(want_enum, got_enum); 53 | } 54 | } 55 | 56 | TEST(SimpleTypesUtilTest, TypeHasNoBaseType) { 57 | { 58 | // artifact type has no base type 59 | SystemTypeExtension extension; 60 | ArtifactType type; 61 | EXPECT_TRUE(absl::IsNotFound(GetSystemTypeExtension(type, extension))); 62 | } 63 | { 64 | // execution type has no base type 65 | SystemTypeExtension extension; 66 | ExecutionType type; 67 | EXPECT_TRUE(absl::IsNotFound(GetSystemTypeExtension(type, extension))); 68 | } 69 | } 70 | 71 | TEST(SimpleTypesUtilTest, BaseTypeIsUnset) { 72 | { 73 | // unset artifact type as base_type 74 | SystemTypeExtension extension; 75 | ArtifactType type; 76 | type.set_base_type(ArtifactType::UNSET); 77 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeExtension(type, extension)); 78 | EXPECT_TRUE(IsUnsetBaseType(extension)); 79 | 80 | ArtifactType::SystemDefinedBaseType got_enum; 81 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeEnum(extension, got_enum)); 82 | EXPECT_THAT(ArtifactType::UNSET, got_enum); 83 | } 84 | { 85 | // unset execution type as base_type 86 | SystemTypeExtension extension; 87 | ExecutionType type; 88 | type.set_base_type(ExecutionType::UNSET); 89 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeExtension(type, extension)); 90 | EXPECT_TRUE(IsUnsetBaseType(extension)); 91 | 92 | ExecutionType::SystemDefinedBaseType got_enum; 93 | ASSERT_EQ(absl::OkStatus(), GetSystemTypeEnum(extension, got_enum)); 94 | EXPECT_EQ(ExecutionType::UNSET, got_enum); 95 | } 96 | } 97 | 98 | } // namespace 99 | } // namespace testing 100 | } // namespace ml_metadata 101 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_metadata_access_object_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | // Test suite for a SqliteMetadataSource based MetadataAccessObject. 16 | 17 | #include 18 | #include 19 | #include 20 | 21 | #include 22 | #include 23 | #include 24 | #include "absl/memory/memory.h" 25 | #include "absl/status/status.h" 26 | #include "ml_metadata/metadata_store/metadata_access_object_factory.h" 27 | #include "ml_metadata/metadata_store/metadata_access_object_test.h" 28 | #include "ml_metadata/metadata_store/metadata_source.h" 29 | #include "ml_metadata/metadata_store/sqlite_metadata_source.h" 30 | #include "ml_metadata/proto/metadata_source.pb.h" 31 | #include "ml_metadata/util/metadata_source_query_config.h" 32 | 33 | namespace ml_metadata { 34 | namespace testing { 35 | 36 | namespace { 37 | 38 | // SqliteMetadataAccessObjectContainer implements MetadataAccessObjectContainer 39 | // to generate and retrieve a MetadataAccessObject based on a 40 | // SqliteMetadataSource. 41 | class SqliteMetadataAccessObjectContainer 42 | : public QueryConfigMetadataAccessObjectContainer { 43 | public: 44 | SqliteMetadataAccessObjectContainer( 45 | std::optional earlier_schema_version = absl::nullopt) 46 | : QueryConfigMetadataAccessObjectContainer( 47 | util::GetSqliteMetadataSourceQueryConfig(), 48 | earlier_schema_version) { 49 | SqliteMetadataSourceConfig config; 50 | metadata_source_ = std::make_unique(config); 51 | CHECK_EQ( 52 | absl::OkStatus(), 53 | CreateMetadataAccessObject( 54 | util::GetSqliteMetadataSourceQueryConfig(), metadata_source_.get(), 55 | earlier_schema_version, &metadata_access_object_)); 56 | } 57 | 58 | ~SqliteMetadataAccessObjectContainer() override = default; 59 | 60 | MetadataSource* GetMetadataSource() override { 61 | return metadata_source_.get(); 62 | } 63 | MetadataAccessObject* GetMetadataAccessObject() override { 64 | return metadata_access_object_.get(); 65 | } 66 | 67 | private: 68 | std::unique_ptr metadata_source_; 69 | std::unique_ptr metadata_access_object_; 70 | }; 71 | 72 | } // namespace 73 | 74 | INSTANTIATE_TEST_SUITE_P( 75 | SqliteMetadataAccessObjectTest, MetadataAccessObjectTest, 76 | ::testing::Values( 77 | []() { 78 | return std::make_unique(); 79 | }, 80 | // TODO(b/257334039) Cleanup after V10+ migration 81 | []() { 82 | return std::make_unique( 83 | /*earlier_schema_version=*/9); 84 | }, 85 | []() { 86 | return std::make_unique( 87 | /*earlier_schema_version=*/8); 88 | }, 89 | []() { 90 | return std::make_unique( 91 | /*earlier_schema_version=*/7); 92 | })); 93 | 94 | } // namespace testing 95 | } // namespace ml_metadata 96 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_metadata_source.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef ML_METADATA_METADATA_STORE_SQLITE_METADATA_SOURCE_H_ 16 | #define ML_METADATA_METADATA_STORE_SQLITE_METADATA_SOURCE_H_ 17 | 18 | #include "absl/status/status.h" 19 | #include "absl/status/statusor.h" 20 | #include "ml_metadata/metadata_store/metadata_source.h" 21 | #include "ml_metadata/proto/metadata_store.pb.h" 22 | #include "sqlite3.h" 23 | 24 | namespace ml_metadata { 25 | 26 | // A MetadataSource based on Sqlite3. By default it uses a in memory Sqlite3 27 | // database, and destroys it when the metadata source is destructed. It can be 28 | // configured via a SqliteMetadataSourceConfig to use physical Sqlite3 and open 29 | // it in read only, read and write, and create if not exists modes. 30 | // This class is thread-unsafe. Multiple objects can be created by using the 31 | // same SqliteMetadataSourceConfig to use the same Sqlite3 database. 32 | class SqliteMetadataSource : public MetadataSource { 33 | public: 34 | explicit SqliteMetadataSource(const SqliteMetadataSourceConfig& config); 35 | ~SqliteMetadataSource() override; 36 | 37 | // Disallow copy and assign. 38 | SqliteMetadataSource(const SqliteMetadataSource&) = delete; 39 | SqliteMetadataSource& operator=(const SqliteMetadataSource&) = delete; 40 | 41 | // Escape strings having single quotes using built-in printf in Sqlite3 C API. 42 | std::string EscapeString(absl::string_view value) const final; 43 | 44 | // Sqlite doesn't escape certain bytes reliably; this is resolved by 45 | // base64 encoding bytes for storage and decoding at access 46 | std::string EncodeBytes(absl::string_view value) const final; 47 | 48 | // DecodeBytes can return an absl::Status if decoding fails 49 | absl::StatusOr DecodeBytes(absl::string_view value) const final; 50 | 51 | private: 52 | // Creates an in memory db. 53 | // If error happens, Returns INTERNAL error. 54 | absl::Status ConnectImpl() final; 55 | 56 | // Closes in memory db. All data stored will be cleaned up. 57 | absl::Status CloseImpl() final; 58 | 59 | // Executes a SQL statement and returns the rows if any. 60 | absl::Status ExecuteQueryImpl(const std::string& query, 61 | RecordSet* results) final; 62 | 63 | // Commits a transaction. 64 | absl::Status CommitImpl() final; 65 | 66 | // Rollbacks a transaction 67 | absl::Status RollbackImpl() final; 68 | 69 | // Begins a transaction 70 | absl::Status BeginImpl() final; 71 | 72 | 73 | // Util methods to execute query. 74 | absl::Status RunStatement(const std::string& query, RecordSet* results); 75 | 76 | // The sqlite3 handle to a database. 77 | sqlite3* db_ = nullptr; 78 | 79 | // A config including connection parameters. 80 | SqliteMetadataSourceConfig config_; 81 | }; 82 | 83 | } // namespace ml_metadata 84 | 85 | #endif // ML_METADATA_METADATA_STORE_SQLITE_METADATA_SOURCE_H_ 86 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_metadata_source_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/sqlite_metadata_source.h" 16 | 17 | #include 18 | 19 | #include 20 | #include 21 | #include "ml_metadata/metadata_store/metadata_source_test_suite.h" 22 | #include "ml_metadata/metadata_store/test_util.h" 23 | 24 | namespace ml_metadata { 25 | namespace testing { 26 | 27 | namespace { 28 | using ml_metadata::testing::EqualsProto; 29 | 30 | class SqliteMetadataSourceContainer : public MetadataSourceContainer { 31 | public: 32 | SqliteMetadataSourceContainer() : MetadataSourceContainer() { 33 | SqliteMetadataSourceConfig config; 34 | metadata_source_ = std::make_unique(config); 35 | } 36 | 37 | explicit SqliteMetadataSourceContainer(SqliteMetadataSourceConfig config) 38 | : MetadataSourceContainer() { 39 | metadata_source_ = std::make_unique(config); 40 | } 41 | 42 | ~SqliteMetadataSourceContainer() override = default; 43 | 44 | MetadataSource* GetMetadataSource() override { 45 | return metadata_source_.get(); 46 | } 47 | 48 | // InitTestSchema creates a new table t1(c1 INT, c2 VARCHAR(255)). 49 | void InitTestSchema() override { 50 | CHECK_EQ(absl::OkStatus(), metadata_source_->Connect()); 51 | CHECK_EQ(absl::OkStatus(), metadata_source_->Begin()); 52 | CHECK_EQ(absl::OkStatus(), 53 | metadata_source_->ExecuteQuery( 54 | "CREATE TABLE t1 (c1 INT, c2 VARCHAR(255));", nullptr)); 55 | CHECK_EQ(absl::OkStatus(), metadata_source_->Commit()); 56 | } 57 | 58 | // InitSchemaAndPopulateRows creates table t1(c1 INT, c2 VARCHAR(255)) and 59 | // adds 3 rows to this table: (1,'v1'), (2,'v2'), (3, 'v3'). 60 | void InitSchemaAndPopulateRows() override { 61 | InitTestSchema(); 62 | CHECK_EQ(absl::OkStatus(), metadata_source_->Begin()); 63 | CHECK_EQ(absl::OkStatus(), metadata_source_->ExecuteQuery( 64 | "INSERT INTO t1 VALUES (1, 'v1')", nullptr)); 65 | CHECK_EQ(absl::OkStatus(), metadata_source_->ExecuteQuery( 66 | "INSERT INTO t1 VALUES (2, 'v2')", nullptr)); 67 | CHECK_EQ(absl::OkStatus(), metadata_source_->ExecuteQuery( 68 | "INSERT INTO t1 VALUES (3, 'v3')", nullptr)); 69 | CHECK_EQ(absl::OkStatus(), metadata_source_->Commit()); 70 | } 71 | 72 | private: 73 | // by default, use a in-memory Sqlite3 metadata source 74 | std::unique_ptr metadata_source_; 75 | }; 76 | 77 | 78 | // Test EscapeString utility method. 79 | TEST(SqliteMetadataSourceExtendedTest, TestEscapeString) { 80 | SqliteMetadataSourceContainer container; 81 | MetadataSource* metadata_source = container.GetMetadataSource(); 82 | EXPECT_EQ(metadata_source->EscapeString("''"), "''''"); 83 | EXPECT_EQ(metadata_source->EscapeString("'\'"), "''''"); 84 | EXPECT_EQ(metadata_source->EscapeString("'\"text\"'"), "''\"text\"''"); 85 | } 86 | 87 | } // namespace 88 | 89 | INSTANTIATE_TEST_SUITE_P( 90 | SqliteMetadataSourceCommonTest, MetadataSourceTestSuite, 91 | ::testing::Values([]() { 92 | return std::make_unique(); 93 | })); 94 | GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(MetadataSourceTestSuite); 95 | } // namespace testing 96 | } // namespace ml_metadata 97 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_metadata_source_util.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/metadata_store/sqlite_metadata_source_util.h" 16 | #include 17 | 18 | #include "ml_metadata/proto/metadata_source.pb.h" 19 | #include "ml_metadata/metadata_store/constants.h" 20 | #include "sqlite3.h" 21 | #include "absl/strings/escaping.h" 22 | #include "absl/status/statusor.h" 23 | #include "absl/strings/string_view.h" 24 | 25 | namespace ml_metadata { 26 | 27 | std::string SqliteEscapeString(absl::string_view value) { 28 | char* buffer = sqlite3_mprintf("%q", value.data()); // NOLINT 29 | std::string result(buffer); 30 | sqlite3_free(buffer); 31 | 32 | return result; 33 | } 34 | 35 | // Encode bytes to workaround escaping difficulties. 36 | std::string SqliteEncodeBytes(absl::string_view value) { 37 | return absl::Base64Escape(value); 38 | } 39 | 40 | absl::StatusOr SqliteDecodeBytes(absl::string_view value) { 41 | absl::StatusOr decoded_or = std::string(); 42 | if (!absl::Base64Unescape(value, &decoded_or.value())) { 43 | decoded_or = absl::InternalError(absl::StrCat( 44 | "Failed to Base64Unescape value '", value, "'")); 45 | } 46 | return decoded_or; 47 | } 48 | 49 | int ConvertSqliteResultsToRecordSet(void* results, int column_num, 50 | char** column_vals, char** column_names) { 51 | // return if queries return no result, e.g., create, insert, update, etc. 52 | if (column_num == 0) return SQLITE_OK; 53 | RecordSet* record_set = static_cast(results); 54 | // ignore the results of the query, if the user passes a nullptr for results. 55 | if (record_set == nullptr) return SQLITE_OK; 56 | // parse the query results to the record_sets (column_names and records) 57 | // as sqlite calls the callback repetitively for each row, it sets the column 58 | // names at the first time. 59 | bool is_column_name_initted = (record_set->column_names_size() == column_num); 60 | if (!is_column_name_initted) record_set->clear_column_names(); 61 | RecordSet::Record* record = record_set->add_records(); 62 | for (int i = 0; i < column_num; i++) { 63 | if (!is_column_name_initted) record_set->add_column_names(column_names[i]); 64 | record->add_values(column_vals[i] ? column_vals[i] 65 | : kMetadataSourceNull.data()); 66 | } 67 | return SQLITE_OK; 68 | } 69 | 70 | } // namespace ml_metadata 71 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_metadata_source_util.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef ML_METADATA_METADATA_STORE_SQLITE_METADATA_SOURCE_UTIL_H_ 16 | #define ML_METADATA_METADATA_STORE_SQLITE_METADATA_SOURCE_UTIL_H_ 17 | 18 | #include 19 | 20 | #include "absl/status/statusor.h" 21 | #include "absl/strings/string_view.h" 22 | 23 | namespace ml_metadata { 24 | 25 | // Escapes strings having single quotes using built-in printf in Sqlite3 C API. 26 | std::string SqliteEscapeString(absl::string_view value); 27 | 28 | // Encode bytes values before binding to queries. 29 | std::string SqliteEncodeBytes(absl::string_view value); 30 | 31 | // Decode bytes values that were encoded with SqliteEncodeBytes 32 | absl::StatusOr SqliteDecodeBytes(absl::string_view value); 33 | 34 | // Converts the query results (`column_vals`) if any to a RecordSet (`results`). 35 | // It is used as a callback of sqlite3_exec. The `results` should be owned by 36 | // the caller of sqlite3_exec. If the given RecordSet (`results`) is nullptr, 37 | // the query result is ignored. 38 | // (see https://www.sqlite.org/c3ref/exec.html for details) 39 | int ConvertSqliteResultsToRecordSet(void* results, int column_num, 40 | char** column_vals, char** column_names); 41 | 42 | } // namespace ml_metadata 43 | 44 | #endif // ML_METADATA_METADATA_STORE_SQLITE_METADATA_SOURCE_UTIL_H_ 45 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_query_config_executor_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | // Test suite for a sqlite query config-based QueryExecutor. 16 | #include 17 | 18 | #include 19 | #include 20 | #include 21 | #include "absl/memory/memory.h" 22 | #include "ml_metadata/metadata_store/query_config_executor.h" 23 | #include "ml_metadata/metadata_store/query_executor.h" 24 | #include "ml_metadata/metadata_store/query_executor_test.h" 25 | #include "ml_metadata/metadata_store/sqlite_metadata_source.h" 26 | #include "ml_metadata/proto/metadata_source.pb.h" 27 | #include "ml_metadata/util/metadata_source_query_config.h" 28 | 29 | namespace ml_metadata { 30 | namespace testing { 31 | 32 | namespace { 33 | // SqliteQueryConfigExecutorContainer implements 34 | // QueryConfigExecutorContainer to generate and retrieve a 35 | // QueryExecutor based on a SqliteMetadataSource. 36 | class SqliteQueryConfigExecutorContainer : public QueryConfigExecutorContainer { 37 | public: 38 | SqliteQueryConfigExecutorContainer() 39 | : QueryConfigExecutorContainer( 40 | util::GetSqliteMetadataSourceQueryConfig()) { 41 | SqliteMetadataSourceConfig config; 42 | metadata_source_ = std::make_unique(config); 43 | if (!metadata_source_->is_connected()) 44 | CHECK_EQ(absl::OkStatus(), metadata_source_->Connect()); 45 | query_executor_ = absl::WrapUnique(new QueryConfigExecutor( 46 | util::GetSqliteMetadataSourceQueryConfig(), metadata_source_.get())); 47 | } 48 | 49 | ~SqliteQueryConfigExecutorContainer() override = default; 50 | 51 | MetadataSource* GetMetadataSource() override { 52 | return metadata_source_.get(); 53 | } 54 | QueryExecutor* GetQueryExecutor() override { return query_executor_.get(); } 55 | 56 | private: 57 | std::unique_ptr metadata_source_; 58 | std::unique_ptr query_executor_; 59 | }; 60 | 61 | } // namespace 62 | 63 | INSTANTIATE_TEST_SUITE_P( 64 | SqliteQueryConfigExecutorTest, QueryExecutorTest, ::testing::Values([]() { 65 | return std::make_unique(); 66 | })); 67 | GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(QueryExecutorTest); 68 | } // namespace testing 69 | } // namespace ml_metadata 70 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/sqlite_rdbms_metadata_access_object_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | // Test suite for a SqliteMetadataSource based RDBMSMetadataAccessObject. 16 | #include 17 | 18 | #include 19 | #include 20 | #include 21 | #include "absl/memory/memory.h" 22 | #include "absl/status/status.h" 23 | #include "ml_metadata/metadata_store/metadata_source.h" 24 | #include "ml_metadata/metadata_store/query_config_executor.h" 25 | #include "ml_metadata/metadata_store/rdbms_metadata_access_object.h" 26 | #include "ml_metadata/metadata_store/rdbms_metadata_access_object_test.h" 27 | #include "ml_metadata/metadata_store/sqlite_metadata_source.h" 28 | #include "ml_metadata/proto/metadata_source.pb.h" 29 | #include "ml_metadata/util/metadata_source_query_config.h" 30 | 31 | namespace ml_metadata { 32 | namespace testing { 33 | 34 | namespace { 35 | 36 | absl::Status CreateRDBMSMetadataAccessObject( 37 | const MetadataSourceQueryConfig& query_config, 38 | MetadataSource* const metadata_source, 39 | std::unique_ptr* result) { 40 | if (!metadata_source->is_connected()) 41 | MLMD_RETURN_IF_ERROR(metadata_source->Connect()); 42 | std::unique_ptr executor = 43 | absl::WrapUnique(new QueryConfigExecutor(query_config, metadata_source)); 44 | *result = 45 | absl::WrapUnique(new RDBMSMetadataAccessObject(std::move(executor))); 46 | return absl::OkStatus(); 47 | } 48 | 49 | // SqliteRDBMSMetadataAccessObjectContainer implements 50 | // RDBMSMetadataAccessObjectContainer to generate and retrieve a 51 | // RDBMSMetadataAccessObject based on a SqliteMetadataSource. 52 | class SqliteRDBMSMetadataAccessObjectContainer 53 | : public QueryConfigRDBMSMetadataAccessObjectContainer { 54 | public: 55 | SqliteRDBMSMetadataAccessObjectContainer() 56 | : QueryConfigRDBMSMetadataAccessObjectContainer( 57 | util::GetSqliteMetadataSourceQueryConfig()) { 58 | SqliteMetadataSourceConfig config; 59 | metadata_source_ = std::make_unique(config); 60 | CHECK_EQ( 61 | absl::OkStatus(), 62 | CreateRDBMSMetadataAccessObject( 63 | util::GetSqliteMetadataSourceQueryConfig(), metadata_source_.get(), 64 | &rmdbs_metadata_access_object_)); 65 | } 66 | 67 | ~SqliteRDBMSMetadataAccessObjectContainer() override = default; 68 | 69 | MetadataSource* GetMetadataSource() override { 70 | return metadata_source_.get(); 71 | } 72 | RDBMSMetadataAccessObject* GetRDBMSMetadataAccessObject() override { 73 | return rmdbs_metadata_access_object_.get(); 74 | } 75 | 76 | private: 77 | std::unique_ptr metadata_source_; 78 | std::unique_ptr rmdbs_metadata_access_object_; 79 | }; 80 | 81 | } // namespace 82 | 83 | INSTANTIATE_TEST_SUITE_P( 84 | SqliteRDBMSMetadataAccessObjectTest, RDBMSMetadataAccessObjectTest, 85 | ::testing::Values([]() { 86 | return std::make_unique(); 87 | })); 88 | GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(RDBMSMetadataAccessObjectTest); 89 | } // namespace testing 90 | } // namespace ml_metadata 91 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/test_mysql_metadata_source_initializer.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_TEST_MYSQL_METADATA_SOURCE_INITIALIZER_H_ 16 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_TEST_MYSQL_METADATA_SOURCE_INITIALIZER_H_ 17 | 18 | #include 19 | 20 | #include "ml_metadata/metadata_store/mysql_metadata_source.h" 21 | 22 | namespace ml_metadata { 23 | namespace testing { 24 | 25 | // TestMySqlMetadataSourceInitializer provides an interface to 26 | // initialize/cleanup a MySqlMetadataSource in tests. 27 | // 28 | // Typical usage: 29 | // 30 | // class FooTest : public ::testing::Test { 31 | // protected: 32 | // void SetUp() override { 33 | // metadata_source_initializer_ = GetTestMySqlMetadataSourceInitializer(); 34 | // metadata_source_ = metadata_source_initializer_->Init( 35 | // TestMysqlMetadataSourceInitializer::ConnectionType::kTcp); 36 | // } 37 | // void TearDown() override { metadata_source_initializer_->Cleanup(); } 38 | // 39 | // TestMySqlMetadataSourceInitializer* metadata_source_initializer_; 40 | // MySqlMetadataSource* metadata_source_; 41 | // }; 42 | class TestMySqlMetadataSourceInitializer { 43 | public: 44 | enum class ConnectionType { kTcp, kSocket }; 45 | 46 | virtual ~TestMySqlMetadataSourceInitializer() = default; 47 | 48 | // Creates or initializes a MySqlMetadataSource. 49 | virtual MySqlMetadataSource* Init(ConnectionType connection_type) = 0; 50 | 51 | // Removes any existing MySqlMetadataSource. 52 | virtual void Cleanup() = 0; 53 | }; 54 | 55 | // Returns a TestMySqlMetadataSourceInitializer to init/cleanup a 56 | // MySqlMetadataSource in tests. 57 | std::unique_ptr 58 | GetTestMySqlMetadataSourceInitializer(); 59 | 60 | } // namespace testing 61 | } // namespace ml_metadata 62 | 63 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_TEST_MYSQL_METADATA_SOURCE_INITIALIZER_H_ 64 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/test_postgresql_metadata_source_initializer.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2023 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | 16 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_TEST_POSTGRESQL_METADATA_SOURCE_INITIALIZER_H_ 17 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_TEST_POSTGRESQL_METADATA_SOURCE_INITIALIZER_H_ 18 | 19 | #include 20 | 21 | #include "ml_metadata/metadata_store/postgresql_metadata_source.h" 22 | 23 | namespace ml_metadata { 24 | namespace testing { 25 | 26 | // Parent class for PostgreSQL database creation and metadata_source 27 | // management. Inherit this class to develop initialization of 28 | // PostgreSQLMetadataSource based on the testing environment. 29 | class TestPostgreSQLMetadataSourceInitializer { 30 | public: 31 | virtual ~TestPostgreSQLMetadataSourceInitializer() = default; 32 | 33 | // Creates or initializes a PostgreSQLMetadataSource. 34 | virtual PostgreSQLMetadataSource* Init() = 0; 35 | 36 | // Removes any existing PostgreSQLMetadataSource. 37 | virtual void Cleanup() = 0; 38 | }; 39 | 40 | // Returns a TestPostgreSQLMetadataSourceInitializer to init/cleanup a 41 | // PostgreSQLMetadataSource in tests. 42 | std::unique_ptr 43 | GetTestPostgreSQLMetadataSourceInitializer(); 44 | 45 | } // namespace testing 46 | } // namespace ml_metadata 47 | 48 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_TEST_POSTGRESQL_METADATA_SOURCE_INITIALIZER_H_ 49 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/test_standalone_mysql_metadata_source_initializer.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include 16 | 17 | #include "gflags/gflags.h" 18 | #include 19 | 20 | #include "absl/memory/memory.h" 21 | #include "ml_metadata/metadata_store/mysql_metadata_source.h" 22 | #include "ml_metadata/metadata_store/test_mysql_metadata_source_initializer.h" 23 | #include "ml_metadata/proto/metadata_store.pb.h" 24 | 25 | DEFINE_string(db_name, "", "Name of MySQL database to connect to"); 26 | DEFINE_string(user_name, "", "MYSQL login id"); 27 | DEFINE_string(password, "", 28 | "Password for user_name. If empty, only MYSQL user ids that " 29 | "don't have a password set are allowed to connect."); 30 | DEFINE_string(host_name, "", "Host name or IP address of the MYSQL server."); 31 | DEFINE_int32(port, 3306, 32 | "TCP port number that the MYSQL server accepts connection on. If " 33 | "not set, the default MYSQL port (3306) is used."); 34 | DEFINE_string(socket, "", "Unix socket file for connecting to MYSQL server."); 35 | 36 | namespace ml_metadata { 37 | namespace testing { 38 | 39 | // Provides a TestMySqlMetadataSourcesInitializer implementation that 40 | // instantiates a MySqlMetadataSource backed by a MYSQL server as specified by 41 | // the flags listed above. 42 | class TestStandaloneMySqlMetadataSourceInitializer 43 | : public TestMySqlMetadataSourceInitializer { 44 | public: 45 | TestStandaloneMySqlMetadataSourceInitializer() = default; 46 | ~TestStandaloneMySqlMetadataSourceInitializer() override = default; 47 | 48 | MySqlMetadataSource* Init(ConnectionType connection_type) override { 49 | MySQLDatabaseConfig config; 50 | config.set_database((FLAGS_db_name)); 51 | config.set_user((FLAGS_user_name)); 52 | config.set_password((FLAGS_password)); 53 | switch (connection_type) { 54 | case ConnectionType::kTcp: 55 | config.set_port((FLAGS_port)); 56 | config.set_host((FLAGS_host_name)); 57 | break; 58 | case ConnectionType::kSocket: 59 | config.set_socket((FLAGS_socket)); 60 | break; 61 | default: 62 | CHECK(false) << "Invalid connection_type: " 63 | << static_cast(connection_type); 64 | } 65 | 66 | metadata_source_ = std::make_unique(config); 67 | CHECK_EQ(absl::OkStatus(), metadata_source_->Connect()); 68 | CHECK_EQ(absl::OkStatus(), metadata_source_->Begin()); 69 | CHECK_EQ(absl::OkStatus(), 70 | metadata_source_->ExecuteQuery( 71 | "DROP DATABASE IF EXISTS " + config.database(), nullptr)); 72 | CHECK_EQ(absl::OkStatus(), metadata_source_->Commit()); 73 | CHECK_EQ(absl::OkStatus(), metadata_source_->Close()); 74 | return metadata_source_.get(); 75 | } 76 | 77 | void Cleanup() override { metadata_source_ = nullptr; } 78 | 79 | private: 80 | std::unique_ptr metadata_source_; 81 | }; 82 | 83 | std::unique_ptr 84 | GetTestMySqlMetadataSourceInitializer() { 85 | return std::make_unique(); 86 | } 87 | 88 | } // namespace testing 89 | } // namespace ml_metadata 90 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/test_util.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | 16 | #ifndef ML_METADATA_METADATA_STORE_TEST_UTIL_H_ 17 | #define ML_METADATA_METADATA_STORE_TEST_UTIL_H_ 18 | 19 | #include 20 | #include 21 | #include 22 | 23 | #include 24 | #include "google/protobuf/text_format.h" 25 | #include "google/protobuf/util/message_differencer.h" 26 | #include 27 | #include 28 | 29 | namespace ml_metadata { 30 | namespace testing { 31 | 32 | template 33 | class ProtoStringMatcher { 34 | public: 35 | explicit ProtoStringMatcher( 36 | const Message& expected, 37 | const std::vector& ignore_fields = {}) 38 | : expected_(expected), ignore_fields_(ignore_fields) {} 39 | 40 | bool MatchAndExplain(const Message& p, 41 | ::testing::MatchResultListener* /* listener */) const; 42 | 43 | void DescribeTo(::std::ostream* os) const { *os << expected_.DebugString(); } 44 | 45 | void DescribeNegationTo(::std::ostream* os) const { 46 | *os << "not equal to expected message: " << expected_.DebugString(); 47 | } 48 | 49 | private: 50 | const Message expected_; 51 | const std::vector& ignore_fields_; 52 | }; 53 | 54 | template 55 | class TupleProtoMatcher : public ::testing::MatcherInterface { 56 | public: 57 | explicit TupleProtoMatcher(const std::vector& ignore_fields) 58 | : ignore_fields_(ignore_fields) {} 59 | 60 | bool MatchAndExplain( 61 | Tuple args, 62 | ::testing::MatchResultListener* /* listener */) const override { 63 | return ProtoStringMatcher(std::get<0>(args), ignore_fields_) 64 | .MatchAndExplain(std::get<1>(args), nullptr); 65 | } 66 | 67 | void DescribeTo(::std::ostream* os) const override { *os << "are equal"; } 68 | 69 | void DescribeNegationTo(::std::ostream* os) const override { 70 | *os << "are not equal"; 71 | } 72 | 73 | private: 74 | const std::vector& ignore_fields_; 75 | }; 76 | 77 | template 78 | bool ProtoStringMatcher::MatchAndExplain( 79 | const Message& p, ::testing::MatchResultListener* /* listener */) const { 80 | google::protobuf::util::MessageDifferencer diff; 81 | for (const std::string& field_name : ignore_fields_) { 82 | const google::protobuf::FieldDescriptor* field_descriptor = 83 | p.descriptor()->FindFieldByName(field_name); 84 | diff.IgnoreField(field_descriptor); 85 | } 86 | return diff.Compare(p, expected_); 87 | } 88 | 89 | // Polymorphic matcher to compare any two protos. 90 | template 91 | inline ::testing::PolymorphicMatcher> EqualsProto( 92 | const Message& x, const std::vector& ignore_fields = {}) { 93 | return ::testing::MakePolymorphicMatcher( 94 | ProtoStringMatcher(x, ignore_fields)); 95 | } 96 | 97 | // A polymorphic matcher for a 2-tuple where first.Equals(second) = true. 98 | template 99 | inline ::testing::Matcher<::testing::tuple> EqualsProto( 100 | const std::vector& ignore_fields = {}) { 101 | return ::testing::MakeMatcher( 102 | new TupleProtoMatcher>( 103 | ignore_fields)); 104 | } 105 | 106 | // Parse input string as a protocol buffer. 107 | template 108 | T ParseTextProtoOrDie(const std::string& input) { 109 | T result; 110 | CHECK(google::protobuf::TextFormat::ParseFromString(input, &result)) 111 | << "Failed to parse: " << input; 112 | return result; 113 | } 114 | 115 | // Checks if the id of a proto matches an expected value. 116 | MATCHER(IdEquals, "") { 117 | auto& lhs = ::testing::get<0>(arg); 118 | auto& rhs = ::testing::get<1>(arg); 119 | return ::testing::Matches(lhs.id())(rhs); 120 | } 121 | 122 | } // namespace testing 123 | } // namespace ml_metadata 124 | 125 | #endif // ML_METADATA_METADATA_STORE_TEST_UTIL_H_ 126 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/test_util_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | 16 | #include "ml_metadata/metadata_store/test_util.h" 17 | 18 | #include 19 | 20 | #include 21 | #include 22 | #include "ml_metadata/proto/metadata_store.pb.h" 23 | 24 | namespace ml_metadata { 25 | namespace testing { 26 | namespace { 27 | 28 | TEST(TestUtilTest, BasicEqualsProto) { 29 | Artifact want_artifact; 30 | want_artifact.set_id(1); 31 | want_artifact.set_uri("some_uri"); 32 | Artifact got_artifact = want_artifact; 33 | EXPECT_THAT(got_artifact, EqualsProto(want_artifact)); 34 | 35 | Artifact other_artifact; 36 | EXPECT_THAT(other_artifact, ::testing::Not(EqualsProto(want_artifact))); 37 | } 38 | 39 | TEST(TestUtilTest, EqualsProtoWithIgnoreFields) { 40 | Artifact want_artifact; 41 | want_artifact.set_id(1); 42 | want_artifact.set_uri("some_uri"); 43 | Artifact empty_artifact; 44 | EXPECT_THAT(empty_artifact, 45 | EqualsProto(want_artifact, /*ignore_fields=*/{"id", "uri"})); 46 | } 47 | 48 | TEST(TestUtilTest, EqualsProtoWithUnorderedPointwise) { 49 | std::vector want_artifacts(2); 50 | want_artifacts[0].set_id(1); 51 | want_artifacts[1].set_uri("2"); 52 | std::vector got_artifacts(2); 53 | got_artifacts[0].set_uri("2"); 54 | got_artifacts[1].set_id(1); 55 | EXPECT_THAT(want_artifacts, ::testing::UnorderedPointwise( 56 | EqualsProto(), got_artifacts)); 57 | got_artifacts[0].set_uri("1"); 58 | EXPECT_THAT(want_artifacts, ::testing::Not(::testing::UnorderedPointwise( 59 | EqualsProto(), got_artifacts))); 60 | EXPECT_THAT( 61 | want_artifacts, 62 | ::testing::UnorderedPointwise( 63 | EqualsProto(/*ignore_fields=*/{"uri"}), got_artifacts)); 64 | } 65 | 66 | 67 | } // namespace 68 | } // namespace testing 69 | } // namespace ml_metadata 70 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/transaction_executor.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | 16 | #include "ml_metadata/metadata_store/transaction_executor.h" 17 | 18 | #include "absl/status/status.h" 19 | #include "ml_metadata/util/return_utils.h" 20 | 21 | namespace ml_metadata { 22 | 23 | absl::Status RdbmsTransactionExecutor::Execute( 24 | const std::function& txn_body, 25 | const TransactionOptions& transaction_options) const { 26 | if (metadata_source_ == nullptr || !metadata_source_->is_connected()) { 27 | return absl::FailedPreconditionError( 28 | "To use ExecuteTransaction, the metadata_source should be created and " 29 | "connected"); 30 | } 31 | 32 | MLMD_RETURN_IF_ERROR(metadata_source_->Begin()); 33 | 34 | absl::Status transaction_status = txn_body(); 35 | if (transaction_status.ok()) { 36 | transaction_status.Update(metadata_source_->Commit()); 37 | } 38 | // Commit may fail as well, if so, we do rollback to allow the caller retry. 39 | if (!transaction_status.ok()) { 40 | transaction_status.Update(metadata_source_->Rollback()); 41 | } 42 | return transaction_status; 43 | } 44 | 45 | } // namespace ml_metadata 46 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/transaction_executor.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | 16 | #ifndef THIRD_PARTY_ML_METADATA_METADATA_STORE_TRANSACTION_EXECUTOR_H_ 17 | #define THIRD_PARTY_ML_METADATA_METADATA_STORE_TRANSACTION_EXECUTOR_H_ 18 | 19 | #include "absl/status/status.h" 20 | #include "ml_metadata/metadata_store/metadata_source.h" 21 | #include "ml_metadata/proto/metadata_store.pb.h" 22 | 23 | namespace ml_metadata { 24 | 25 | // Pure virtual interface for MetadataStore to execute a transaction. 26 | // 27 | // Example usage: 28 | // TransactionExecutor* txn_executor; 29 | // ... 30 | // txn_executor->Execute( 31 | // [&metadata_access_object]() -> absl::Status { 32 | // return metadata_access_object->InitMetadataSource(); 33 | // }); 34 | class TransactionExecutor { 35 | public: 36 | virtual ~TransactionExecutor() = default; 37 | 38 | // Runs txn_body and return the transaction status. 39 | virtual absl::Status Execute( 40 | const std::function& txn_body, 41 | const TransactionOptions& transaction_options = TransactionOptions()) 42 | const = 0; 43 | }; 44 | 45 | // An implementation of TransactionExecutor. 46 | // It contains a method to execute the transaction body and tries to commit 47 | // the execution result in the database by using Begin/Commit/Rollback 48 | // methods in MetadataSource. 49 | class RdbmsTransactionExecutor : public TransactionExecutor { 50 | public: 51 | explicit RdbmsTransactionExecutor(MetadataSource* metadata_source) 52 | : metadata_source_(metadata_source) {} 53 | ~RdbmsTransactionExecutor() override = default; 54 | 55 | // Tries to commit the execution result of txn_body. 56 | // When the txn_body returns OK, it calls Commit, otherwise it calls Rollback. 57 | // 58 | // Returns FAILED_PRECONDITION if metadata_source is null or not connected. 59 | // Returns detailed internal errors of transaction, i.e. 60 | // Begin, Rollback and Commit. 61 | absl::Status Execute(const std::function& txn_body, 62 | const TransactionOptions& transaction_options = 63 | TransactionOptions()) const override; 64 | 65 | private: 66 | // The MetadataSource which has the connection to a database. 67 | // It also supports other database primitves like Commit and Abort. 68 | // Not owned by this class. 69 | MetadataSource* metadata_source_; 70 | }; 71 | 72 | } // namespace ml_metadata 73 | 74 | #endif // THIRD_PARTY_ML_METADATA_METADATA_STORE_TRANSACTION_EXECUTOR_H_ 75 | -------------------------------------------------------------------------------- /ml_metadata/metadata_store/types.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef ML_METADATA_METADATA_STORE_TYPES_H_ 16 | #define ML_METADATA_METADATA_STORE_TYPES_H_ 17 | 18 | #include 19 | 20 | #ifndef HAS_GLOBAL_STRING 21 | 22 | using std::string; 23 | 24 | #endif 25 | 26 | #endif // ML_METADATA_METADATA_STORE_TYPES_H_ 27 | -------------------------------------------------------------------------------- /ml_metadata/move_generated_files.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright 2020 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | # Convenience binary to build MLMD from source. 17 | # Put wrapped c++ files in place 18 | # Should be run in the directory containing WORKSPACE file. (workspace root) 19 | 20 | function _is_macos() { 21 | [[ "$(uname -s | tr 'A-Z' 'a-z')" =~ darwin ]] 22 | } 23 | 24 | function mlmd::move_generated_files() { 25 | set -eux 26 | 27 | # Newer bazel does not create bazel-genfiles any more (see 28 | # https://github.com/bazelbuild/bazel/issues/6761). It's merged with 29 | # bazel-bin after Bazel 0.25.0. Latest bazel doesn't provide bazel-genfiles 30 | # symlink anymore, so if the symlink directory doesnt' exist, use bazel-bin. 31 | local bazel_genfiles 32 | if [[ -d "${BUILD_WORKSPACE_DIRECTORY}/bazel-genfiles" ]]; then 33 | bazel_genfiles="bazel-genfiles" 34 | else 35 | bazel_genfiles="bazel-bin" 36 | fi 37 | 38 | cp -f ${BUILD_WORKSPACE_DIRECTORY}/${bazel_genfiles}/ml_metadata/proto/metadata_store_pb2.py \ 39 | ${BUILD_WORKSPACE_DIRECTORY}/ml_metadata/proto 40 | cp -f ${BUILD_WORKSPACE_DIRECTORY}/${bazel_genfiles}/ml_metadata/proto/metadata_store_service_pb2.py \ 41 | ${BUILD_WORKSPACE_DIRECTORY}/ml_metadata/proto 42 | cp -f ${BUILD_WORKSPACE_DIRECTORY}/${bazel_genfiles}/ml_metadata/proto/metadata_store_service_pb2_grpc.py \ 43 | ${BUILD_WORKSPACE_DIRECTORY}/ml_metadata/proto 44 | cp -f ${BUILD_WORKSPACE_DIRECTORY}/${bazel_genfiles}/ml_metadata/simple_types/proto/simple_types_pb2.py \ 45 | ${BUILD_WORKSPACE_DIRECTORY}/ml_metadata/simple_types/proto 46 | 47 | 48 | MLMD_EXTENSION="ml_metadata/metadata_store/pywrap/metadata_store_extension.so" 49 | cp -f ${BUILD_WORKSPACE_DIRECTORY}/bazel-bin/${MLMD_EXTENSION} \ 50 | ${BUILD_WORKSPACE_DIRECTORY}/${MLMD_EXTENSION} 51 | 52 | chmod +w "${BUILD_WORKSPACE_DIRECTORY}/${MLMD_EXTENSION}" 53 | } 54 | 55 | mlmd::move_generated_files 56 | -------------------------------------------------------------------------------- /ml_metadata/mysql_configure.bzl: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """ 16 | Pulls mariadb-connector-c github repo and builds it. 17 | 18 | Clients can then depend on `@libmysqlclient` to use the MYSQL C API. 19 | """ 20 | 21 | load("@bazel_tools//tools/build_defs/repo:git.bzl", "new_git_repository") 22 | 23 | def mysql_configure(): 24 | new_git_repository( 25 | name = "libmysqlclient", 26 | build_file = "//ml_metadata:libmysqlclient.BUILD", 27 | workspace_file = "//ml_metadata:libmysqlclient.WORKSPACE", 28 | remote = "https://github.com/MariaDB/mariadb-connector-c.git", 29 | tag = "v3.0.8-release", 30 | patches = ["//ml_metadata/third_party:libmysqlclient.patch"], 31 | ) 32 | -------------------------------------------------------------------------------- /ml_metadata/postgresql.WORKSPACE: -------------------------------------------------------------------------------- 1 | workspace(name = "postgresql") 2 | -------------------------------------------------------------------------------- /ml_metadata/proto/BUILD: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | load("//ml_metadata:ml_metadata.bzl", "ml_metadata_proto_library", "ml_metadata_proto_library_go", "ml_metadata_proto_library_py") 16 | 17 | package(default_visibility = ["//visibility:public"]) 18 | 19 | licenses(["notice"]) 20 | 21 | ml_metadata_proto_library( 22 | name = "metadata_store_proto", 23 | srcs = ["metadata_store.proto"], 24 | deps = ["@com_google_protobuf//:cc_wkt_protos"], 25 | ) 26 | 27 | ml_metadata_proto_library( 28 | name = "metadata_store_service_proto", 29 | srcs = ["metadata_store_service.proto"], 30 | has_services = 1, 31 | cc_grpc_version = 1, 32 | deps = [":metadata_store_proto"], 33 | ) 34 | 35 | ml_metadata_proto_library_py( 36 | name = "metadata_store_py_pb2", 37 | srcs = ["metadata_store.proto"], 38 | api_version = 2, 39 | proto_library = "metadata_store_proto", 40 | ) 41 | 42 | ml_metadata_proto_library_py( 43 | name = "metadata_store_service_py_pb2", 44 | srcs = ["metadata_store_service.proto"], 45 | api_version = 2, 46 | oss_deps = [":metadata_store_py_pb2"], 47 | proto_library = "metadata_store_service_proto", 48 | use_grpc_plugin = True, 49 | ) 50 | 51 | ml_metadata_proto_library_go( 52 | name = "metadata_store_go_proto", 53 | srcs = ["metadata_store.proto"], 54 | importpath = "ml_metadata/proto/metadata_store_go_proto", 55 | deps = [":metadata_store_proto"], 56 | ) 57 | 58 | ml_metadata_proto_library_go( 59 | name = "metadata_store_service_go_proto", 60 | srcs = ["metadata_store_service.proto"], 61 | cc_proto_deps = [":metadata_store_proto"], 62 | gen_oss_grpc = True, 63 | go_proto_deps = [":metadata_store_go_proto"], 64 | importpath = "ml_metadata/proto/metadata_store_service_go_proto", 65 | deps = [":metadata_store_service_proto"], 66 | ) 67 | 68 | ml_metadata_proto_library( 69 | name = "metadata_source_proto", 70 | srcs = ["metadata_source.proto"], 71 | deps = [ 72 | "@com_google_protobuf//:cc_wkt_protos", 73 | ], 74 | ) 75 | -------------------------------------------------------------------------------- /ml_metadata/proto/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """ML Metadata proto module.""" 15 | 16 | # Connection configurations for different deployment. 17 | from ml_metadata.proto import ( 18 | metadata_store_service_pb2, 19 | metadata_store_service_pb2_grpc, 20 | ) 21 | 22 | # ML Metadata core data model concepts. 23 | from ml_metadata.proto.metadata_store_pb2 import ( 24 | Artifact, 25 | ArtifactType, 26 | Association, 27 | Attribution, 28 | ConnectionConfig, 29 | Context, 30 | ContextType, 31 | Event, 32 | Execution, 33 | ExecutionType, 34 | FakeDatabaseConfig, 35 | MetadataStoreClientConfig, 36 | MySQLDatabaseConfig, 37 | ParentContext, 38 | SqliteMetadataSourceConfig, 39 | ) 40 | 41 | del metadata_store_service_pb2 42 | del metadata_store_service_pb2_grpc 43 | 44 | Artifact.__doc__ = """ 45 | An artifact represents an input or an output of individual steps in a ML 46 | workflow, e.g., a trained model, an input dataset, and evaluation metrics. 47 | """ 48 | 49 | Execution.__doc__ = """ 50 | An execution describes a component run or a step in an ML workflow along with 51 | its runtime parameters, e.g., a Trainer run, a data transformation step. 52 | """ 53 | 54 | Context.__doc__ = """ 55 | A context defines a group of artifacts and/or executions. It captures the 56 | commonalities of the grouped entities. For example, a project context that 57 | contains all artifacts and executions of a ML workflow, an experiment that 58 | includes the modified trainers and evaluation metrics. 59 | """ 60 | 61 | Event.__doc__ = """ 62 | An event records the relationship between artifacts and executions. When an 63 | execution happens, events record every artifact that was used by the execution, 64 | and every artifact that was produced. These records allow for lineage tracking 65 | throughout a ML workflow. 66 | """ 67 | 68 | Attribution.__doc__ = """ 69 | An attribution represents the relationship between artifacts and contexts. 70 | """ 71 | 72 | Association.__doc__ = """ 73 | An association represents the relationship between executions and contexts. 74 | """ 75 | 76 | ParentContext.__doc__ = """ 77 | A parental context represents the relationship between contexts. 78 | """ 79 | 80 | ArtifactType.__doc__ = """ 81 | A user defined type about a collection of artifacts and their properties that 82 | are stored in the metadata store. 83 | """ 84 | 85 | ExecutionType.__doc__ = """ 86 | A user defined type about a collection of executions and their properties that 87 | are stored in the metadata store. 88 | """ 89 | 90 | ContextType.__doc__ = """ 91 | A user defined type about a collection of contexts and their properties that 92 | are stored in the metadata store. 93 | """ 94 | 95 | ConnectionConfig.__doc__ = """ 96 | A connection configuration specifying the persistent backend to be used with 97 | MLMD. 98 | """ 99 | 100 | FakeDatabaseConfig.__doc__ = """ 101 | An in-memory database configuration for testing purpose. 102 | """ 103 | 104 | MySQLDatabaseConfig.__doc__ = """ 105 | A connection configuration to use a MySQL db instance as a MLMD backend. 106 | """ 107 | 108 | SqliteMetadataSourceConfig.__doc__ = """ 109 | A connection configuration to use a Sqlite db file as a MLMD backend. 110 | """ 111 | 112 | MetadataStoreClientConfig.__doc__ = """ 113 | A connection configuration to use a MLMD server as the persistent backend. 114 | """ 115 | 116 | __all__ = [ 117 | "ConnectionConfig", 118 | "MetadataStoreClientConfig", 119 | "Artifact", 120 | "Execution", 121 | "Context", 122 | "Event", 123 | "Attribution", 124 | "Association", 125 | "ParentContext", 126 | "ArtifactType", 127 | "ExecutionType", 128 | "ContextType", 129 | "FakeDatabaseConfig", 130 | "MySQLDatabaseConfig", 131 | "SqliteMetadataSourceConfig", 132 | ] 133 | -------------------------------------------------------------------------------- /ml_metadata/proto/testing/BUILD: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # This mock protobuffer is used to unit test PROTO value storage. 16 | # The mock_proto cc_library provides a utility function for building the proto 17 | 18 | load("//ml_metadata:ml_metadata.bzl", "ml_metadata_proto_library") 19 | 20 | package(default_visibility = ["//visibility:public"]) 21 | 22 | licenses(["notice"]) 23 | 24 | ml_metadata_proto_library( 25 | name = "mock_proto", 26 | srcs = ["mock.proto"], 27 | ) 28 | -------------------------------------------------------------------------------- /ml_metadata/proto/testing/mock.proto: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | syntax = "proto2"; 16 | 17 | package ml_metadata.testing; 18 | 19 | option java_multiple_files = true; 20 | 21 | 22 | message MockProto { 23 | optional string string_value = 1; 24 | optional double double_value = 2; 25 | } 26 | -------------------------------------------------------------------------------- /ml_metadata/py.typed: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /ml_metadata/query/BUILD: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | load( 16 | "//ml_metadata:ml_metadata.bzl", 17 | "ml_metadata_cc_test", 18 | ) 19 | 20 | licenses(["notice"]) 21 | 22 | cc_library( 23 | name = "filter_query_ast_resolver", 24 | srcs = ["filter_query_ast_resolver.cc"], 25 | hdrs = ["filter_query_ast_resolver.h"], 26 | visibility = ["//ml_metadata:__subpackages__"], 27 | deps = [ 28 | "@com_google_absl//absl/container:flat_hash_set", 29 | "@com_google_absl//absl/status", 30 | "@com_google_absl//absl/status:statusor", 31 | "@com_google_absl//absl/strings", 32 | "//ml_metadata/proto:metadata_store_proto", 33 | "//ml_metadata/util:return_utils", 34 | "@com_googlesource_code_re2//:re2", 35 | "@com_google_zetasql//zetasql/public:analyzer", 36 | "@com_google_zetasql//zetasql/public:simple_catalog", 37 | ], 38 | ) 39 | 40 | ml_metadata_cc_test( 41 | name = "filter_query_ast_resolver_test", 42 | size = "small", 43 | srcs = ["filter_query_ast_resolver_test.cc"], 44 | deps = [ 45 | ":filter_query_ast_resolver", 46 | "@com_google_googletest//:gtest_main", 47 | "@com_google_absl//absl/status", 48 | "//ml_metadata/metadata_store:test_util", 49 | "//ml_metadata/proto:metadata_store_proto", 50 | ], 51 | ) 52 | 53 | cc_library( 54 | name = "filter_query_builder", 55 | srcs = ["filter_query_builder.cc"], 56 | hdrs = ["filter_query_builder.h"], 57 | visibility = ["//ml_metadata:__subpackages__"], 58 | deps = [ 59 | "@com_google_absl//absl/container:btree", 60 | "@com_google_absl//absl/status", 61 | "@com_google_absl//absl/strings", 62 | "//ml_metadata/metadata_store:constants", 63 | "//ml_metadata/proto:metadata_store_proto", 64 | "@com_google_glog//:glog", 65 | "@com_google_zetasql//zetasql/public:strings", 66 | "@com_google_zetasql//zetasql/resolved_ast:sql_builder", 67 | ], 68 | ) 69 | 70 | ml_metadata_cc_test( 71 | name = "filter_query_builder_test", 72 | size = "small", 73 | srcs = ["filter_query_builder_test.cc"], 74 | deps = [ 75 | ":filter_query_ast_resolver", 76 | ":filter_query_builder", 77 | "@com_google_googletest//:gtest_main", 78 | "@com_google_absl//absl/status", 79 | "//ml_metadata/metadata_store:test_util", 80 | "//ml_metadata/proto:metadata_store_proto", 81 | ], 82 | ) 83 | -------------------------------------------------------------------------------- /ml_metadata/query/filter_query_ast_resolver.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef ML_METADATA_GOOGLE_QUERY_FILTER_QUERY_AST_RESOLVER_H 16 | #define ML_METADATA_GOOGLE_QUERY_FILTER_QUERY_AST_RESOLVER_H 17 | 18 | #include "zetasql/public/analyzer.h" 19 | #include "zetasql/public/simple_catalog.h" 20 | #include "absl/status/status.h" 21 | #include "re2/re2.h" 22 | 23 | namespace ml_metadata { 24 | 25 | // FilterQueryAstResolver parses the MLMD filtering query string and generates 26 | // an AST via ZetaSQL analyzer. It can be instantiated with MLMD nodes types: 27 | // Artifact, Execution and Context. 28 | template 29 | class FilterQueryAstResolver { 30 | public: 31 | // Constructs an AST resolver for a given query_string. Always use a new AST 32 | // resolver instance to process a query string, as the temporary types for 33 | // each query string is not sharable. 34 | explicit FilterQueryAstResolver(const std::string& query_string); 35 | 36 | // Not copyable or movable 37 | FilterQueryAstResolver(const FilterQueryAstResolver&) = delete; 38 | FilterQueryAstResolver& operator=(const FilterQueryAstResolver&) = delete; 39 | 40 | // Processes the user query string and generates an ast. 41 | // It reads the user query string, creates a list of zetasql types 42 | // and validates the query. 43 | // Returns InvalidArgument if the given query is not a boolean expression. 44 | // Returns InvalidArgument errors if the query string contains symbols that 45 | // is not supported by the filter query syntax and semantics. 46 | absl::Status Resolve(); 47 | 48 | // Returns the parsed AST of the giving query string. The AST is owned by 49 | // the instance of FilterQueryAstResolver and it is immutable. 50 | // The AST can be used with FilterQueryBuilder to generate actual query. 51 | // Returns nullptr if the Resolve() failed with errors. 52 | const zetasql::ResolvedExpr* GetAst(); 53 | 54 | private: 55 | // The user query. 56 | const std::string raw_query_; 57 | // The type factory to keep the temporary types used in the query. 58 | zetasql::TypeFactory type_factory_; 59 | // The analyzer options that bridge the alias and the type. 60 | zetasql::AnalyzerOptions analyzer_opts_; 61 | // The simple catalog used by the analyzer for resolving functions. 62 | zetasql::SimpleCatalog catalog_; 63 | // The output that holds the resolved query. 64 | std::unique_ptr output_; 65 | }; 66 | 67 | } // namespace ml_metadata 68 | 69 | #endif // ML_METADATA_GOOGLE_QUERY_FILTER_QUERY_AST_RESOLVER_H 70 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/BUILD: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | package(default_visibility = ["//ml_metadata:__subpackages__"]) 16 | 17 | licenses(["notice"]) 18 | 19 | cc_library( 20 | name = "simple_types_constants", 21 | srcs = [ 22 | "simple_types_constants.cc", 23 | ], 24 | hdrs = [ 25 | "simple_types_constants.h", 26 | ], 27 | deps = [ 28 | "@com_google_absl//absl/base:core_headers", 29 | "@com_google_absl//absl/strings", 30 | ], 31 | ) 32 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/proto/BUILD: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | load("//ml_metadata:ml_metadata.bzl", "ml_metadata_proto_library", "ml_metadata_proto_library_py") 15 | 16 | package(default_visibility = ["//ml_metadata:__subpackages__"]) 17 | 18 | licenses(["notice"]) 19 | 20 | ml_metadata_proto_library( 21 | name = "simple_types_proto", 22 | srcs = ["simple_types.proto"], 23 | deps = [ 24 | "//ml_metadata/proto:metadata_store_proto", 25 | ], 26 | ) 27 | 28 | ml_metadata_proto_library_py( 29 | name = "simple_types_py_pb2", 30 | srcs = ["simple_types.proto"], 31 | api_version = 2, 32 | oss_deps = ["//ml_metadata/proto:metadata_store_py_pb2"], 33 | proto_library = "simple_types_proto", 34 | ) 35 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/proto/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/proto/simple_types.proto: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | syntax = "proto3"; 16 | 17 | package ml_metadata; 18 | 19 | import "ml_metadata/proto/metadata_store.proto"; 20 | 21 | // A bundle of ml-metadata types to describe artifacts, executions and contexts 22 | // in general ML pipelines. 23 | // The details of the data model is described in go/mlmd. 24 | // ml-metadata provides a predefined bundle defined in simple_types_constants.h. 25 | message SimpleTypes { 26 | // A list of artifact types. 27 | repeated ArtifactType artifact_types = 1; 28 | 29 | // A list of execution types. 30 | repeated ExecutionType execution_types = 2; 31 | 32 | // A list of context types. 33 | repeated ContextType context_types = 3; 34 | } 35 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/simple_types_constants.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/simple_types/simple_types_constants.h" 16 | 17 | #include 18 | 19 | #include "absl/strings/string_view.h" 20 | 21 | namespace ml_metadata { 22 | 23 | // LINT.IfChange 24 | const absl::string_view kSimpleTypes = R"pb( 25 | artifact_types { name: "mlmd.Dataset" } 26 | artifact_types { name: "mlmd.Model" } 27 | artifact_types { name: "mlmd.Metrics" } 28 | artifact_types { name: "mlmd.Statistics" } 29 | execution_types { name: "mlmd.Train" } 30 | execution_types { name: "mlmd.Transform" } 31 | execution_types { name: "mlmd.Process" } 32 | execution_types { name: "mlmd.Evaluate" } 33 | execution_types { name: "mlmd.Deploy" } 34 | )pb"; 35 | // LINT.ThenChange(../proto/metadata_store.proto, 36 | // ../metadata_store/simple_types_util.cc) 37 | 38 | const std::array kSimpleTypeNames = { 39 | "mlmd.Dataset", "mlmd.Model", "mlmd.Metrics", 40 | "mlmd.Statistics", "mlmd.Train", "mlmd.Transform", 41 | "mlmd.Process", "mlmd.Evaluate", "mlmd.Deploy"}; 42 | 43 | } // namespace ml_metadata 44 | -------------------------------------------------------------------------------- /ml_metadata/simple_types/simple_types_constants.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2020 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_SIMPLE_TYPES_SIMPLE_TYPES_CONSTANTS_H_ 16 | #define THIRD_PARTY_ML_METADATA_SIMPLE_TYPES_SIMPLE_TYPES_CONSTANTS_H_ 17 | 18 | #include 19 | 20 | #include "absl/base/attributes.h" 21 | #include "absl/strings/string_view.h" 22 | 23 | namespace ml_metadata { 24 | 25 | static const int kNumSimpleTypes = 9; 26 | 27 | // The in-memory encoding of SimpleTypes proto. 28 | ABSL_CONST_INIT extern const absl::string_view kSimpleTypes; 29 | 30 | // A list of simple type names. 31 | ABSL_CONST_INIT extern const std::array 32 | kSimpleTypeNames; 33 | 34 | } // namespace ml_metadata 35 | 36 | #endif // THIRD_PARTY_ML_METADATA_SIMPLE_TYPES_SIMPLE_TYPES_CONSTANTS_H_ 37 | -------------------------------------------------------------------------------- /ml_metadata/third_party/BUILD: -------------------------------------------------------------------------------- 1 | licenses(["notice"]) 2 | -------------------------------------------------------------------------------- /ml_metadata/third_party/glog.BUILD: -------------------------------------------------------------------------------- 1 | # Glog has it's own BUILD file, but it enables gflags by default. 2 | # Rather than add another dependency, this BUILD file is exactly 3 | # the same but without gflags. 4 | 5 | licenses(["notice"]) 6 | 7 | load(":bazel/glog.bzl", "glog_library") 8 | 9 | glog_library(with_gflags = 0) 10 | -------------------------------------------------------------------------------- /ml_metadata/third_party/grpc.patch: -------------------------------------------------------------------------------- 1 | diff --git bazel/grpc_deps.bzl bazel/grpc_deps.bzl 2 | index ef9b2bea94..0be413d48f 100644 3 | --- bazel/grpc_deps.bzl 4 | +++ bazel/grpc_deps.bzl 5 | @@ -122,7 +122,7 @@ def grpc_deps(): 6 | 7 | native.bind( 8 | name = "re2", 9 | - actual = "@com_googlesource_code_re2//:re2", 10 | + actual = "@com_github_google_re2//:re2", 11 | ) 12 | 13 | native.bind( 14 | -------------------------------------------------------------------------------- /ml_metadata/third_party/libmysqlclient.patch: -------------------------------------------------------------------------------- 1 | --- cmake/ConnectorName.cmake 2 | +++ cmake/ConnectorName.cmake 3 | @@ -22,7 +22,7 @@ IF(CMAKE_SYSTEM_NAME MATCHES "Windows") 4 | SET(MACHINE_NAME "x64") 5 | ELSE() 6 | SET(MACHINE_NAME "32") 7 | - END() 8 | + ENDIF() 9 | ENDIF() 10 | 11 | SET(product_name "mysql-connector-c-${CPACK_PACKAGE_VERSION}-${PLATFORM_NAME}${CONCAT_SIGN}${MACHINE_NAME}") 12 | -------------------------------------------------------------------------------- /ml_metadata/third_party/rules_foreign_cc.patch: -------------------------------------------------------------------------------- 1 | diff --git foreign_cc/private/cc_toolchain_util.bzl foreign_cc/private/cc_toolchain_util.bzl 2 | index 1bd872d..7a7880d 100644 3 | --- foreign_cc/private/cc_toolchain_util.bzl 4 | +++ foreign_cc/private/cc_toolchain_util.bzl 5 | @@ -46,6 +46,14 @@ FOREIGN_CC_DISABLED_FEATURES = [ 6 | "thin_lto", 7 | ] 8 | 9 | +def _maybe_convert_gcc(gcc_path): 10 | + gcc_path_steps = gcc_path.split("/") 11 | + if (gcc_path_steps[-1] == "gcc"): 12 | + gcc_path_steps[-1] = "g++" 13 | + return "/".join(gcc_path_steps) 14 | + else: 15 | + return gcc_path 16 | + 17 | def _to_list(element): 18 | if element == None: 19 | return [] 20 | @@ -270,10 +278,10 @@ def get_tools_info(ctx): 21 | feature_configuration = feature_configuration, 22 | action_name = ACTION_NAMES.c_compile, 23 | ), 24 | - cxx = cc_common.get_tool_for_action( 25 | + cxx = _maybe_convert_gcc(cc_common.get_tool_for_action( 26 | feature_configuration = feature_configuration, 27 | action_name = ACTION_NAMES.cpp_compile, 28 | - ), 29 | + )), 30 | cxx_linker_static = cc_common.get_tool_for_action( 31 | feature_configuration = feature_configuration, 32 | action_name = ACTION_NAMES.cpp_link_static_library, 33 | -------------------------------------------------------------------------------- /ml_metadata/third_party/six.BUILD: -------------------------------------------------------------------------------- 1 | # This file is copied from https://github.com/abseil/abseil-py/blob/master/third_party/six.BUILD. 2 | # Description: 3 | # Six provides simple utilities for wrapping over differences between Python 2 4 | # and Python 3. 5 | 6 | licenses(["notice"]) # MIT 7 | 8 | exports_files(["LICENSE"]) 9 | 10 | py_library( 11 | name = "six", 12 | srcs = ["six.py"], 13 | srcs_version = "PY2AND3", 14 | visibility = ["//visibility:public"], 15 | ) 16 | -------------------------------------------------------------------------------- /ml_metadata/third_party/sqlite.BUILD: -------------------------------------------------------------------------------- 1 | # Description: 2 | # sqlite3 is a serverless SQL RDBMS. 3 | 4 | licenses(["unencumbered"]) # Public Domain 5 | 6 | config_setting( 7 | name = "macos", 8 | constraint_values = [ 9 | "@platforms//os:osx", 10 | ], 11 | ) 12 | 13 | SQLITE_COPTS = [ 14 | "-DSQLITE_ENABLE_JSON1", 15 | "-DHAVE_DECL_STRERROR_R=1", 16 | "-DHAVE_STDINT_H=1", 17 | "-DHAVE_INTTYPES_H=1", 18 | "-D_FILE_OFFSET_BITS=64", 19 | "-D_REENTRANT=1", 20 | ] + select({ 21 | ":macos": [ 22 | "-Os", 23 | "-DHAVE_GMTIME_R=1", 24 | "-DHAVE_LOCALTIME_R=1", 25 | "-DHAVE_USLEEP=1", 26 | ], 27 | "//conditions:default": [ 28 | "-Os", 29 | "-DHAVE_FDATASYNC=1", 30 | "-DHAVE_GMTIME_R=1", 31 | "-DHAVE_LOCALTIME_R=1", 32 | "-DHAVE_POSIX_FALLOCATE=1", 33 | "-DHAVE_USLEEP=1", 34 | ], 35 | }) 36 | 37 | # Production build of SQLite library that's baked into TensorFlow. 38 | cc_library( 39 | name = "org_sqlite", 40 | srcs = ["sqlite3.c"], 41 | hdrs = [ 42 | "sqlite3.h", 43 | "sqlite3ext.h", 44 | ], 45 | copts = SQLITE_COPTS, 46 | defines = [ 47 | # This gets rid of the bloat of deprecated functionality. It 48 | # needs to be listed here instead of copts because it's actually 49 | # referenced in the sqlite3.h file. 50 | "SQLITE_OMIT_DEPRECATED", 51 | ], 52 | linkopts = select({ 53 | "//conditions:default": [ 54 | "-ldl", 55 | "-lpthread", 56 | ], 57 | }), 58 | visibility = ["//visibility:public"], 59 | ) 60 | 61 | # This is a Copybara sync helper for Google. 62 | py_library( 63 | name = "python", 64 | srcs_version = "PY3", 65 | visibility = ["//visibility:public"], 66 | ) 67 | -------------------------------------------------------------------------------- /ml_metadata/tools/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google LLC. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /ml_metadata/tools/dev_debug/Dockerfile: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # Choose docker base image. 16 | FROM ubuntu:20.04 as builder 17 | 18 | # Install required dependencies. 19 | RUN export DEBIAN_FRONTEND=noninteractive && \ 20 | apt-get update -y --option Acquire::Retries=3 && \ 21 | apt-get install --no-install-recommends -y -q --option Acquire::Retries=3 \ 22 | build-essential \ 23 | clang \ 24 | cmake \ 25 | make \ 26 | musl-dev \ 27 | openssl \ 28 | curl \ 29 | ca-certificates \ 30 | unzip \ 31 | software-properties-common \ 32 | git \ 33 | python3-distutils \ 34 | python-is-python3 \ 35 | python3-dev 36 | 37 | # Set up Bazel 6.1.0 38 | ENV BAZEL_VERSION 6.1.0 39 | WORKDIR / 40 | RUN mkdir /bazel && \ 41 | cd /bazel && \ 42 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 43 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -o /bazel/LICENSE.txt https://raw.githubusercontent.com/bazelbuild/bazel/master/LICENSE && \ 44 | chmod +x bazel-*.sh && \ 45 | ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 46 | cd / && \ 47 | rm -f /bazel/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh 48 | 49 | ADD . /dev_debug 50 | WORKDIR /dev_debug 51 | -------------------------------------------------------------------------------- /ml_metadata/tools/docker_build/Dockerfile.manylinux2010: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # Dockerfile for building a manylinux2010 MLMD wheel. 16 | 17 | # This docker image is essentially pypa/manylinux2010 + bazel. 18 | FROM gcr.io/tfx-oss-public/manylinux2014-bazel:bazel-6.1.0 19 | WORKDIR /build 20 | CMD ["ml_metadata/tools/docker_build/build_manylinux.sh"] 21 | -------------------------------------------------------------------------------- /ml_metadata/tools/docker_build/build_manylinux.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright 2019 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | # This script is expected to run in the docker container defined in 17 | # Dockerfile.manylinux2010 18 | # Assumptions: 19 | # - CentOS environment. 20 | # - devtoolset-8 is installed. 21 | # - $PWD is MLMD's project root. 22 | # - Python of different versions are installed at /opt/python/. 23 | # - patchelf, zip, bazel is installed and is in $PATH. 24 | 25 | WORKING_DIR=$PWD 26 | 27 | function setup_environment() { 28 | source scl_source enable devtoolset-8 29 | source scl_source enable rh-python38 30 | if [[ -z "${PYTHON_VERSION}" ]]; then 31 | echo "Must set PYTHON_VERSION env to 39|310|311|"; exit 1; 32 | fi 33 | # Bazel will use PYTHON_BIN_PATH to determine the right python library. 34 | if [[ "${PYTHON_VERSION}" == 39 ]]; then 35 | PYTHON_DIR=/opt/python/cp39-cp39 36 | elif [[ "${PYTHON_VERSION}" == 310 ]]; then 37 | PYTHON_DIR=/opt/python/cp310-cp310 38 | elif [[ "${PYTHON_VERSION}" == 311 ]]; then 39 | PYTHON_DIR=/opt/python/cp311-cp311 40 | else 41 | echo "Must set PYTHON_VERSION env to 39|310|311|"; exit 1; 42 | fi 43 | export PIP_BIN="${PYTHON_DIR}"/bin/pip 44 | export PYTHON_BIN_PATH="${PYTHON_DIR}"/bin/python 45 | echo "PYTHON_BIN_PATH=${PYTHON_BIN_PATH}" 46 | ${PIP_BIN} install --upgrade pip 47 | ${PIP_BIN} install --upgrade wheel 48 | ${PIP_BIN} install numpy 49 | # Auditwheel does not have a python2 version and auditwheel is just a binary. 50 | ${PIP_BIN} install auditwheel 51 | sudo ln -s ${PYTHON_BIN_PATH} /usr/bin/python3 52 | } 53 | 54 | function build_wheel() { 55 | rm -rf dist 56 | "${PYTHON_BIN_PATH}" setup.py bdist_wheel 57 | } 58 | 59 | function stamp_wheel() { 60 | WHEEL_PATH="$(ls "$PWD"/dist/*.whl)" 61 | WHEEL_DIR=$(dirname "${WHEEL_PATH}") 62 | auditwheel repair --plat manylinux2014_x86_64 -w "${WHEEL_DIR}" "${WHEEL_PATH}" 63 | rm "${WHEEL_PATH}" 64 | } 65 | 66 | setup_environment 67 | set -e 68 | set -x 69 | build_wheel 70 | stamp_wheel 71 | 72 | set +e 73 | set +x 74 | -------------------------------------------------------------------------------- /ml_metadata/tools/docker_server/Dockerfile: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # TODO(b/195701120) Introduces l.gcr.io/google/bazel:5.3.0 when it is available 16 | # and makes sure that it uses ubuntu 20.04 as base image. Currently the lastest 17 | # version only supports bazel 6.1.0. 18 | FROM ubuntu:20.04 as builder 19 | 20 | RUN export DEBIAN_FRONTEND=noninteractive && \ 21 | apt-get update -y --option Acquire::Retries=3 && \ 22 | apt-get install --no-install-recommends -y -q --option Acquire::Retries=3 \ 23 | build-essential \ 24 | clang \ 25 | cmake \ 26 | make \ 27 | musl-dev \ 28 | openssl \ 29 | curl \ 30 | ca-certificates \ 31 | unzip \ 32 | software-properties-common \ 33 | git \ 34 | python3-distutils \ 35 | python-is-python3 \ 36 | python3-dev 37 | 38 | # Set up Bazel 6.1.0 39 | ENV BAZEL_VERSION 6.1.0 40 | WORKDIR / 41 | RUN mkdir /bazel && \ 42 | cd /bazel && \ 43 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 44 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -o /bazel/LICENSE.txt https://raw.githubusercontent.com/bazelbuild/bazel/master/LICENSE && \ 45 | chmod +x bazel-*.sh && \ 46 | ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 47 | cd / && \ 48 | rm -f /bazel/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh 49 | 50 | ADD . /mlmd-src 51 | WORKDIR /mlmd-src 52 | 53 | # "-std=c++17" is needed in order to build with ZetaSQL. 54 | RUN bazel build -c opt --action_env=PATH \ 55 | --define=grpc_no_ares=true \ 56 | //ml_metadata/metadata_store:metadata_store_server --cxxopt="-std=c++17" 57 | 58 | # copying libmysqlclient source onto THIRD_PARTY folder. 59 | RUN mkdir -p /mlmd-src/third_party 60 | RUN cp -RL /mlmd-src/bazel-mlmd-src/external/libmysqlclient /mlmd-src/third_party/mariadb-connector-c 61 | 62 | FROM ubuntu:20.04 63 | 64 | COPY --from=builder /mlmd-src/bazel-bin/ml_metadata/metadata_store/metadata_store_server /bin/metadata_store_server 65 | COPY --from=builder /mlmd-src/third_party /mlmd-src/third_party 66 | 67 | ENV GRPC_PORT "8080" 68 | ENV METADATA_STORE_SERVER_CONFIG_FILE "" 69 | 70 | # Introduces tzdata package here to avoid LoadTimeZone check failed error in the metadata store server. 71 | RUN export DEBIAN_FRONTEND=noninteractive && \ 72 | apt-get update -y --option Acquire::Retries=3 && \ 73 | apt-get install --no-install-recommends -y -q --option Acquire::Retries=3 \ 74 | tzdata 75 | 76 | ENTRYPOINT \ 77 | "/bin/metadata_store_server" \ 78 | "--grpc_port=${GRPC_PORT}" \ 79 | "--metadata_store_server_config_file=${METADATA_STORE_SERVER_CONFIG_FILE}" 80 | -------------------------------------------------------------------------------- /ml_metadata/tools/docker_server/README.md: -------------------------------------------------------------------------------- 1 | Files for building the [Docker](http://www.docker.com) image for running the 2 | MLMD gRPC server. 3 | 4 | To build a docker image, run 5 | `./ml_metadata/tools/docker_server/build_docker_image.sh` under root directory 6 | of github checkout. 7 | -------------------------------------------------------------------------------- /ml_metadata/tools/docker_server/build_docker_image.sh: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # Convenience script to build docker image for MLMD gRPC server. 16 | 17 | set -u -x 18 | 19 | DOCKER_IMAGE_REPO=${DOCKER_IMAGE_REPO:-"gcr.io/tfx-oss-public/ml_metadata_store_server"} 20 | DOCKER_IMAGE_TAG=${DOCKER_IMAGE_TAG:-"latest"} 21 | DOCKER_FILE=${DOCKER_FILE:-"Dockerfile"} 22 | 23 | # Run docker build command. 24 | docker build -t ${DOCKER_IMAGE_REPO}:${DOCKER_IMAGE_TAG} -f ml_metadata/tools/docker_server/${DOCKER_FILE} . 25 | -------------------------------------------------------------------------------- /ml_metadata/util/BUILD: -------------------------------------------------------------------------------- 1 | # Utilities for ML Metadata 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | cc_library( 8 | name = "metadata_source_query_config", 9 | srcs = ["metadata_source_query_config.cc"], 10 | hdrs = ["metadata_source_query_config.h"], 11 | deps = [ 12 | "@com_google_protobuf//:protobuf", 13 | "@com_google_absl//absl/strings", 14 | "//ml_metadata/proto:metadata_source_proto", 15 | "@com_google_glog//:glog", 16 | ], 17 | ) 18 | 19 | cc_library( 20 | name = "return_utils", 21 | hdrs = ["return_utils.h"], 22 | ) 23 | 24 | cc_library( 25 | name = "struct_utils", 26 | srcs = ["struct_utils.cc"], 27 | hdrs = ["struct_utils.h"], 28 | deps = [ 29 | "@com_google_protobuf//:protobuf", 30 | "@com_google_absl//absl/status", 31 | "@com_google_absl//absl/strings", 32 | "@com_google_protobuf//:cc_wkt_protos", 33 | ], 34 | ) 35 | 36 | cc_test( 37 | name = "struct_utils_test", 38 | srcs = ["struct_utils_test.cc"], 39 | deps = [ 40 | ":struct_utils", 41 | "@com_google_protobuf//:protobuf", 42 | "@com_google_googletest//:gtest_main", 43 | "@com_google_absl//absl/status", 44 | "@com_google_absl//absl/strings", 45 | "//ml_metadata/metadata_store:test_util", 46 | "@com_google_glog//:glog", 47 | "@com_google_protobuf//:cc_wkt_protos", 48 | ], 49 | ) 50 | 51 | cc_test( 52 | name = "metadata_source_query_config_test", 53 | srcs = ["metadata_source_query_config_test.cc"], 54 | deps = [ 55 | ":metadata_source_query_config", 56 | "@com_google_googletest//:gtest_main", 57 | ], 58 | ) 59 | 60 | cc_library( 61 | name = "field_mask_utils", 62 | srcs = ["field_mask_utils.cc"], 63 | hdrs = ["field_mask_utils.h"], 64 | deps = [ 65 | "@com_google_absl//absl/container:flat_hash_set", 66 | "@com_google_absl//absl/status", 67 | "@com_google_absl//absl/status:statusor", 68 | "@com_google_absl//absl/strings", 69 | "//ml_metadata/proto:metadata_store_proto", 70 | "@com_google_protobuf//:protobuf", 71 | ], 72 | ) 73 | 74 | cc_test( 75 | name = "field_mask_utils_test", 76 | srcs = ["field_mask_utils_test.cc"], 77 | deps = [ 78 | ":field_mask_utils", 79 | "@com_google_googletest//:gtest_main", 80 | "@com_google_absl//absl/container:flat_hash_set", 81 | "@com_google_absl//absl/status", 82 | "@com_google_absl//absl/status:statusor", 83 | "@com_google_absl//absl/strings", 84 | "//ml_metadata/proto:metadata_store_proto", 85 | "@com_google_protobuf//:protobuf", 86 | ], 87 | ) 88 | 89 | cc_library( 90 | name = "record_parsing_utils", 91 | srcs = ["record_parsing_utils.cc"], 92 | hdrs = ["record_parsing_utils.h"], 93 | deps = [ 94 | ":return_utils", 95 | 96 | "@com_google_absl//absl/container:flat_hash_set", 97 | "@com_google_absl//absl/status", 98 | "@com_google_absl//absl/strings", 99 | "//ml_metadata/metadata_store:constants", 100 | "//ml_metadata/proto:metadata_source_proto", 101 | "//ml_metadata/proto:metadata_store_proto", 102 | "@com_google_glog//:glog", 103 | ], 104 | ) 105 | 106 | cc_test( 107 | name = "record_parsing_utils_test", 108 | srcs = ["record_parsing_utils_test.cc"], 109 | deps = [ 110 | ":record_parsing_utils", 111 | "@com_google_googletest//:gtest_main", 112 | "@com_google_absl//absl/status", 113 | "//ml_metadata/metadata_store:test_util", 114 | "//ml_metadata/proto:metadata_store_proto", 115 | ], 116 | ) 117 | -------------------------------------------------------------------------------- /ml_metadata/util/field_mask_utils.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2023 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/util/field_mask_utils.h" 16 | 17 | #include 18 | 19 | #include "google/protobuf/field_mask.pb.h" 20 | #include "absl/container/flat_hash_set.h" 21 | #include "absl/status/status.h" 22 | #include "absl/status/statusor.h" 23 | #include "absl/strings/match.h" 24 | #include "absl/strings/str_cat.h" 25 | #include "absl/strings/string_view.h" 26 | #include "absl/strings/strip.h" 27 | #include "ml_metadata/proto/metadata_store.pb.h" 28 | #include "google/protobuf/descriptor.h" 29 | 30 | namespace ml_metadata { 31 | 32 | namespace { 33 | constexpr absl::string_view kPropertiesFieldPathPrefix = "properties"; 34 | constexpr absl::string_view kCustomPropertiesFieldPathPrefix = 35 | "custom_properties"; 36 | } // namespace 37 | 38 | absl::StatusOr> GetPropertyNamesFromMask( 39 | const google::protobuf::FieldMask& mask, const bool is_custom_properties) { 40 | absl::flat_hash_set property_names; 41 | absl::string_view prefix = is_custom_properties 42 | ? kCustomPropertiesFieldPathPrefix 43 | : kPropertiesFieldPathPrefix; 44 | for (absl::string_view path : mask.paths()) { 45 | if (path == prefix) { 46 | return absl::InternalError( 47 | absl::StrCat("Cannot split property names from ", prefix, 48 | " if no property name is specified for ", prefix, 49 | " field. Please update ", prefix, " field as a whole.")); 50 | } 51 | if (absl::StartsWith(path, prefix)) { 52 | property_names.insert(absl::StripPrefix(path, absl::StrCat(prefix, "."))); 53 | } 54 | } 55 | return property_names; 56 | } 57 | 58 | absl::StatusOr> 59 | GetPropertyNamesFromMaskOrUnionOfProperties( 60 | const google::protobuf::FieldMask& mask, const bool is_custom_properties, 61 | const google::protobuf::Map& curr_properties, 62 | const google::protobuf::Map& prev_properties) { 63 | absl::flat_hash_set property_names_in_properties_union; 64 | for (const auto& property : curr_properties) { 65 | property_names_in_properties_union.insert(property.first); 66 | } 67 | for (const auto& property : prev_properties) { 68 | property_names_in_properties_union.insert(property.first); 69 | } 70 | if (mask.paths().empty()) { 71 | return property_names_in_properties_union; 72 | } 73 | 74 | absl::StatusOr> 75 | property_names_in_mask_or_internal_error = 76 | GetPropertyNamesFromMask(mask, is_custom_properties); 77 | if (property_names_in_mask_or_internal_error.status().code() == 78 | absl::StatusCode::kInternal) { 79 | return property_names_in_properties_union; 80 | } 81 | 82 | return property_names_in_mask_or_internal_error; 83 | } 84 | 85 | absl::StatusOr GetFieldsSubMaskFromMask( 86 | const google::protobuf::FieldMask& mask, 87 | const google::protobuf::Descriptor* descriptor) { 88 | if (descriptor == nullptr) { 89 | return absl::InvalidArgumentError("`descriptor` cannot be null."); 90 | } 91 | google::protobuf::FieldMask submask; 92 | 93 | std::copy_if(mask.paths().begin(), mask.paths().end(), 94 | google::protobuf::RepeatedFieldBackInserter(submask.mutable_paths()), 95 | [descriptor](absl::string_view p) { 96 | return !absl::StartsWith(p, kPropertiesFieldPathPrefix) && 97 | !absl::StartsWith(p, 98 | kCustomPropertiesFieldPathPrefix) && 99 | descriptor->FindFieldByName(p.data()) != nullptr; 100 | }); 101 | return submask; 102 | } 103 | } // namespace ml_metadata 104 | -------------------------------------------------------------------------------- /ml_metadata/util/field_mask_utils.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2023 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_UTIL_FIELD_MASK_UTILS_H_ 16 | #define THIRD_PARTY_ML_METADATA_UTIL_FIELD_MASK_UTILS_H_ 17 | 18 | #include "google/protobuf/field_mask.pb.h" 19 | #include "absl/container/flat_hash_set.h" 20 | #include "absl/status/statusor.h" 21 | #include "absl/strings/string_view.h" 22 | #include "ml_metadata/proto/metadata_store.pb.h" 23 | #include "google/protobuf/descriptor.h" 24 | 25 | namespace ml_metadata { 26 | 27 | // Gets names of `properties` or `custom_properties` from `mask`. 28 | // For example, if `mask` contains {"properties.x", "properties.y", 29 | // "external_id"} and `is_custom_properties` is false, then the splited results 30 | // will be: {"x", "y"}. 31 | // Returns INTERNAL error if `mask` contains the field path for 32 | // "properties" without specifying the names in it, e.g. {"properties"}, and 33 | // `is_custom_properties` is false. 34 | // Returns INTERNAL error if `mask` contains the field path for 35 | // "custom_properties" without specifying the names in it , 36 | // e.g. {"custom_properties"}, and `is_custom_properties` is true. 37 | absl::StatusOr> GetPropertyNamesFromMask( 38 | const google::protobuf::FieldMask& mask, bool is_custom_properties); 39 | 40 | // Gets names of `properties` or `custom_properties` from `mask` or the union of 41 | // `curr_properties` and `prev_properties` if `mask` is empty. 42 | // For example: 43 | // If `mask` contains {"properties.x", "properties.y", "external_id"} and 44 | // `is_custom_properties` is false, then the result will contain: {"x", "y"} 45 | // If `mask` contains {"external_id"} (does not contain property names), 46 | // then the result will be empty. 47 | // If `mask` is empty, the the result will contain the union of property names 48 | // from both `curr_properties` and `prev_properties`. 49 | // Returns INTERNAL error if any error occurs. 50 | absl::StatusOr> 51 | GetPropertyNamesFromMaskOrUnionOfProperties( 52 | const google::protobuf::FieldMask& mask, const bool is_custom_properties, 53 | const google::protobuf::Map& curr_properties, 54 | const google::protobuf::Map& prev_properties); 55 | 56 | // Gets a valid submask only containing paths of fields from `mask`. 57 | // For example, if `mask` contains {"properties.x", "properties.y", 58 | // "external_id", "invalid_filed_path"} , then the submask will be: 59 | // {"external_id"}. 60 | // Returns INVALID_ARGUMENT error if `descriptor` is a null-pointer. 61 | absl::StatusOr GetFieldsSubMaskFromMask( 62 | const google::protobuf::FieldMask& mask, 63 | const google::protobuf::Descriptor* descriptor); 64 | 65 | } // namespace ml_metadata 66 | 67 | #endif // THIRD_PARTY_ML_METADATA_UTIL_FIELD_MASK_UTILS_H_ 68 | -------------------------------------------------------------------------------- /ml_metadata/util/metadata_source_query_config.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef ML_METADATA_UTIL_METADATA_SOURCE_QUERY_CONFIG_H_ 16 | #define ML_METADATA_UTIL_METADATA_SOURCE_QUERY_CONFIG_H_ 17 | 18 | #include "ml_metadata/proto/metadata_source.pb.h" 19 | 20 | namespace ml_metadata { 21 | namespace util { 22 | 23 | // Gets the MetadataSourceQueryConfig for MySQLMetadataSource. 24 | MetadataSourceQueryConfig GetMySqlMetadataSourceQueryConfig(); 25 | 26 | // Gets the MetadataSourceQueryConfig for SQLiteMetadataSource. 27 | MetadataSourceQueryConfig GetSqliteMetadataSourceQueryConfig(); 28 | 29 | // Gets the MetadataSourceQueryConfig for FakeMetadataSource. 30 | MetadataSourceQueryConfig GetFakeMetadataSourceQueryConfig(); 31 | 32 | // Gets the MetadataSourceQueryConfig for PostgreSQLMetadataSource. 33 | MetadataSourceQueryConfig GetPostgreSQLMetadataSourceQueryConfig(); 34 | 35 | 36 | } // namespace util 37 | } // namespace ml_metadata 38 | #endif // ML_METADATA_UTIL_METADATA_SOURCE_QUERY_CONFIG_H_ 39 | -------------------------------------------------------------------------------- /ml_metadata/util/metadata_source_query_config_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/util/metadata_source_query_config.h" 16 | 17 | #include 18 | 19 | namespace ml_metadata { 20 | namespace util { 21 | namespace { 22 | 23 | TEST(MetadataSourceQueryConfig, GetMySqlMetadataSourceQueryConfig) { 24 | const MetadataSourceQueryConfig config = GetMySqlMetadataSourceQueryConfig(); 25 | EXPECT_EQ(config.metadata_source_type(), MYSQL_METADATA_SOURCE); 26 | } 27 | 28 | TEST(MetadataSourceQueryConfig, GetSqliteMetadataSourceQueryConfig) { 29 | const MetadataSourceQueryConfig config = GetSqliteMetadataSourceQueryConfig(); 30 | EXPECT_EQ(config.metadata_source_type(), SQLITE_METADATA_SOURCE); 31 | } 32 | 33 | 34 | } // namespace 35 | } // namespace util 36 | } // namespace ml_metadata 37 | -------------------------------------------------------------------------------- /ml_metadata/util/return_utils.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_UTIL_RETURN_UTILS_H_ 16 | #define THIRD_PARTY_ML_METADATA_UTIL_RETURN_UTILS_H_ 17 | 18 | namespace ml_metadata { 19 | 20 | // For propagating errors when calling a function. 21 | #define MLMD_RETURN_IF_ERROR(...) \ 22 | do { \ 23 | absl::Status _status = (__VA_ARGS__); \ 24 | if (!_status.ok()) return _status; \ 25 | } while (0) 26 | 27 | #define MLMD_RETURN_WITH_CONTEXT_IF_ERROR(expr, ...) \ 28 | do { \ 29 | absl::Status _status = (expr); \ 30 | if (!_status.ok()) { \ 31 | return absl::Status(_status.code(), \ 32 | absl::StrCat(__VA_ARGS__, _status.message())); \ 33 | } \ 34 | } while (0) 35 | 36 | #define MLMD_STATUS_MACROS_CONCAT_NAME(x, y) \ 37 | MLMD_STATUS_MACROS_CONCAT_IMPL(x, y) 38 | 39 | #define MLMD_STATUS_MACROS_CONCAT_IMPL(x, y) x##y 40 | 41 | // Executes an expression `rexpr` that returns an `absl::StatusOr`. On OK, 42 | // moves its value into the variable defined by `lhs`, otherwise returns 43 | // from the current function. 44 | #define MLMD_ASSIGN_OR_RETURN(lhs, rexpr) \ 45 | MLMD_ASSIGN_OR_RETURN_IMPL( \ 46 | MLMD_STATUS_MACROS_CONCAT_NAME(_status_or_value, __COUNTER__), lhs, \ 47 | rexpr) 48 | 49 | #define MLMD_ASSIGN_OR_RETURN_IMPL(statusor, lhs, rexpr) \ 50 | auto statusor = (rexpr); \ 51 | if (!statusor.ok()) { \ 52 | return statusor.status(); \ 53 | } \ 54 | lhs = std::move(statusor.value()) 55 | 56 | } // namespace ml_metadata 57 | 58 | #endif // THIRD_PARTY_ML_METADATA_UTIL_RETURN_UTILS_H_ 59 | -------------------------------------------------------------------------------- /ml_metadata/util/struct_utils.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/util/struct_utils.h" 16 | 17 | #include "google/protobuf/struct.pb.h" 18 | #include "absl/status/status.h" 19 | #include "absl/strings/escaping.h" 20 | #include "absl/strings/match.h" 21 | #include "absl/strings/str_cat.h" 22 | #include "absl/strings/string_view.h" 23 | #include "absl/strings/strip.h" 24 | 25 | namespace ml_metadata { 26 | namespace { 27 | 28 | // TODO(b/178235219): Remove this once we migrate to schema v7 and add a column 29 | // for storing serialized `Struct` values. 30 | constexpr absl::string_view kSerializedStructPrefix = "mlmd-struct::"; 31 | } 32 | 33 | bool IsStructSerializedString(absl::string_view serialized_value) { 34 | return absl::StartsWith(serialized_value, kSerializedStructPrefix); 35 | } 36 | 37 | std::string StructToString(const google::protobuf::Struct& struct_value) { 38 | return absl::StrCat(kSerializedStructPrefix, 39 | absl::Base64Escape(struct_value.SerializeAsString())); 40 | } 41 | 42 | absl::Status StringToStruct(absl::string_view serialized_value, 43 | google::protobuf::Struct& struct_value) { 44 | if (!absl::StartsWith(serialized_value, kSerializedStructPrefix)) { 45 | return absl::InvalidArgumentError( 46 | absl::StrCat("Not a valid serialized `Struct`: ", serialized_value)); 47 | } 48 | std::string unescaped; 49 | 50 | if (!absl::Base64Unescape( 51 | absl::StripPrefix(serialized_value, kSerializedStructPrefix), 52 | &unescaped)) { 53 | return absl::InvalidArgumentError(absl::StrCat( 54 | "Unable to parse serialized `google.protobuf.Struct` value: ", 55 | unescaped)); 56 | } 57 | 58 | if (!struct_value.ParseFromString(unescaped)) { 59 | return absl::InvalidArgumentError(absl::StrCat( 60 | "Unable to parse serialized `google.protobuf.Struct` value: ", 61 | unescaped)); 62 | } 63 | 64 | return absl::OkStatus(); 65 | } 66 | 67 | } // namespace ml_metadata 68 | -------------------------------------------------------------------------------- /ml_metadata/util/struct_utils.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2021 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #ifndef THIRD_PARTY_ML_METADATA_UTIL_STRUCT_UTILS_H_ 16 | #define THIRD_PARTY_ML_METADATA_UTIL_STRUCT_UTILS_H_ 17 | 18 | #include 19 | 20 | #include "google/protobuf/struct.pb.h" 21 | #include "absl/status/status.h" 22 | #include "absl/strings/string_view.h" 23 | 24 | namespace ml_metadata { 25 | 26 | // Returns true if |serialized_value| is a `google.protobuf.Struct` serialized 27 | // for storage. 28 | bool IsStructSerializedString(absl::string_view serialized_value); 29 | 30 | // Returns a serialized string representing |struct_value| suitable for storage 31 | // in MLMD's Property table. 32 | std::string StructToString(const google::protobuf::Struct& struct_value); 33 | 34 | // Returns the corresponding `google.protobuf.Struct` deserialized from 35 | // |serialized_value| in |struct_value|. 36 | // Returns INVALID_ARGUMENT if the value in |serialized_value| does not look 37 | // like a valid string serialized using |StructToString| above. 38 | absl::Status StringToStruct(absl::string_view serialized_value, 39 | google::protobuf::Struct& struct_value); 40 | 41 | } // namespace ml_metadata 42 | 43 | #endif // THIRD_PARTY_ML_METADATA_UTIL_STRUCT_UTILS_H_ 44 | -------------------------------------------------------------------------------- /ml_metadata/util/struct_utils_test.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2019 Google LLC 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | https://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | ==============================================================================*/ 15 | #include "ml_metadata/util/struct_utils.h" 16 | 17 | #include 18 | #include "google/protobuf/struct.pb.h" 19 | #include 20 | #include 21 | #include "absl/status/status.h" 22 | #include "absl/strings/str_cat.h" 23 | #include "ml_metadata/metadata_store/test_util.h" 24 | 25 | namespace ml_metadata { 26 | namespace { 27 | 28 | using ::ml_metadata::testing::EqualsProto; 29 | 30 | TEST(StructUtils, IsStructStringWorks) { 31 | EXPECT_TRUE( 32 | IsStructSerializedString("mlmd-struct::")); 33 | EXPECT_FALSE( 34 | IsStructSerializedString("mlmd-struct:")); 35 | EXPECT_FALSE( 36 | IsStructSerializedString("mlmd--struct::")); 37 | EXPECT_FALSE(IsStructSerializedString("")); 38 | } 39 | 40 | TEST(StructUtils, SerializeAndDeserializeStructWorks) { 41 | const auto struct_value = 42 | ml_metadata::testing::ParseTextProtoOrDie(R"pb( 43 | fields { 44 | key: "json number" 45 | value { number_value: 1234 } 46 | } 47 | fields { 48 | key: "json object" 49 | value { 50 | struct_value { 51 | fields { 52 | key: "nested json key" 53 | value { string_value: "string value" } 54 | } 55 | } 56 | } 57 | } 58 | )pb"); 59 | 60 | const std::string serialized_value = StructToString(struct_value); 61 | google::protobuf::Struct got_value; 62 | ASSERT_EQ(absl::OkStatus(), StringToStruct(serialized_value, got_value)); 63 | EXPECT_THAT(got_value, EqualsProto(struct_value)); 64 | } 65 | 66 | TEST(StructUtils, DeserializeInvalidPrefix) { 67 | google::protobuf::Struct got_value; 68 | const absl::Status status = StringToStruct("mlmd-struct--", got_value); 69 | EXPECT_TRUE(absl::IsInvalidArgument(status)); 70 | } 71 | 72 | TEST(StructUtils, DeserializeNonBase64) { 73 | google::protobuf::Struct got_value; 74 | const absl::Status status = StringToStruct("mlmd-struct::garbage", got_value); 75 | EXPECT_TRUE(absl::IsInvalidArgument(status)); 76 | } 77 | 78 | TEST(StructUtils, DeserializeNonStructSerializedString) { 79 | const google::protobuf::Struct stored_value = 80 | ml_metadata::testing::ParseTextProtoOrDie(R"pb( 81 | fields { 82 | key: "json number" 83 | value { number_value: 1234 } 84 | })pb"); 85 | google::protobuf::Struct got_value; 86 | const absl::Status status = StringToStruct( 87 | absl::StrCat("mlmd-struct::", stored_value.SerializeAsString()), 88 | got_value); 89 | EXPECT_TRUE(absl::IsInvalidArgument(status)); 90 | } 91 | 92 | } // namespace 93 | } // namespace ml_metadata 94 | -------------------------------------------------------------------------------- /ml_metadata/version.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Contains the version string of ML Metadata.""" 16 | 17 | # Note that setup.py uses this version. 18 | __version__ = '1.18.0.dev' 19 | -------------------------------------------------------------------------------- /ml_metadata/workspace.bzl: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """ML METADATA Data Validation external dependencies that can be loaded in WORKSPACE files. 16 | """ 17 | 18 | load("//ml_metadata:mysql_configure.bzl", "mysql_configure") 19 | 20 | def ml_metadata_workspace(): 21 | """All ML Metadata external dependencies.""" 22 | mysql_configure() 23 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | [build-system] 16 | requires = [ 17 | "setuptools", 18 | "wheel", 19 | # Required for using org_tensorflow bazel repository. 20 | "numpy~=1.22.0", 21 | ] 22 | 23 | [tool.pytest.ini_options] 24 | testpaths = ["ml_metadata/metadata_store/"] 25 | python_files = "*_test.py" 26 | 27 | [tool.ruff] 28 | line-length = 88 29 | 30 | [tool.ruff.lint] 31 | select = [ 32 | # pycodestyle 33 | "E", 34 | "W", 35 | # Pyflakes 36 | "F", 37 | # pyupgrade 38 | "UP", 39 | # flake8-bugbear 40 | "B", 41 | # flake8-simplify 42 | "SIM", 43 | # isort 44 | "I", 45 | # pep8 naming 46 | "N", 47 | # pydocstyle 48 | "D", 49 | # annotations 50 | "ANN", 51 | # debugger 52 | "T10", 53 | # flake8-pytest 54 | "PT", 55 | # flake8-return 56 | "RET", 57 | # flake8-unused-arguments 58 | "ARG", 59 | # flake8-fixme 60 | "FIX", 61 | # flake8-eradicate 62 | "ERA", 63 | # pandas-vet 64 | "PD", 65 | # numpy-specific rules 66 | "NPY", 67 | ] 68 | ignore = [ 69 | "D104", # Missing docstring in public package 70 | "D100", # Missing docstring in public module 71 | "D211", # No blank line before class 72 | "D213", # Multiline summary second line 73 | "PD901", # Avoid using 'df' for pandas dataframes. Perfectly fine in functions with limited scope 74 | "ANN201", # Missing return type annotation for public function (makes no sense for NoneType return types...) 75 | "ANN101", # Missing type annotation for `self` 76 | "ANN204", # Missing return type annotation for special method 77 | "ANN002", # Missing type annotation for `*args` 78 | "ANN003", # Missing type annotation for `**kwargs` 79 | "D105", # Missing docstring in magic method 80 | "D203", # 1 blank line before after class docstring 81 | "D204", # 1 blank line required after class docstring 82 | "D413", # 1 blank line after parameters 83 | "SIM108", # Simplify if/else to one line; not always clearer 84 | "D206", # Advised to disable by ruff-format 85 | "E501", # Advised to disable by ruff-format 86 | "W191", # Advised to disable by ruff-format 87 | "N802", # Function name should be lowercase; unittest uses mixed case 88 | "D107", # Missing docstring in `__init__` 89 | 90 | # These are issues which remain to be fixed 91 | "D102", # Missing docstring in public method 92 | "D103", # Missing docstring in public function 93 | "D401", # First line of docstring should be in imperative mood 94 | "D404", # First word of the docstring should not be "This" 95 | "ANN001", # Missing type annotation for function argument 96 | "ANN202", # Missing return type for private function 97 | "B024", # is an abstract base class, but it has no abstract methods 98 | "FIX002", # Line contains TODO, consider resolving the issue 99 | "ANN401", # Dynamically typed expressions (typing.Any) are disallowed 100 | "UP008", # Use `super()` instead of `super(__class__, self)` 101 | "SIM102", # Use a single `if` statement instead of nested `if` statements 102 | ] 103 | 104 | [tool.ruff.lint.per-file-ignores] 105 | "__init__.py" = [ 106 | "F401", # Unused import 107 | ] 108 | "*_test.py" = [ 109 | "ANN001", # Type annotations aren't needed for tests; these are fixtures or parametrizations 110 | "PT009", # Use a regular `assert` instead of a unittest-style `assertEqual` 111 | "PT027", # Use `pytest.raises` instead of unittest-style `assertRaisesRegex` 112 | 113 | # Missing docstrings; probably want to fill these out for tests. For now, we just disable 114 | "D101", # Missing docstring in public class 115 | "D102", # Missing docstring in public method 116 | 117 | ] 118 | 119 | [tool.ruff.lint.pydocstyle] 120 | convention = "google" 121 | 122 | [tool.isort] 123 | profile = "black" 124 | -------------------------------------------------------------------------------- /requirements-docs.txt: -------------------------------------------------------------------------------- 1 | mkdocs 2 | mkdocs-material 3 | mkdocstrings[python] 4 | griffe-inherited-docstrings 5 | mkdocs-autorefs 6 | black 7 | mkdocs-jupyter 8 | mkdocs-caption 9 | -------------------------------------------------------------------------------- /test_constraints.txt: -------------------------------------------------------------------------------- 1 | # test_constraints.txt 2 | # This file specifies the constraints for the test environment of mlmd. 3 | # Unlike library dependency which aims to specify the widest version range 4 | # possible, it is okay to specify exact version here. 5 | # 6 | # test_constraints.txt file is similar to requirements.txt except it does not 7 | # tell to really "install" the specified target; it only specifies the version 8 | # constraint if it is installed either directly or transitively by the 9 | # dependencies. 10 | 11 | # TODO(b/281693330): Remove this after upgrading OpenSSL version in our test 12 | # environment. 13 | urllib3<2 14 | --------------------------------------------------------------------------------