├── .github ├── config.yml └── workflows │ ├── checks.yml │ └── upgrade_automation.yml ├── .gitignore ├── .golangci.yml ├── .goreleaser.yml ├── CODEOWNERS ├── CODE_OF_CONDUCT.md ├── Dockerfile ├── LICENSE ├── Makefile ├── NOTICE ├── README.md ├── boilerplate ├── flyte │ ├── code_of_conduct │ │ ├── CODE_OF_CONDUCT.md │ │ ├── README.rst │ │ └── update.sh │ ├── docker_build │ │ ├── Makefile │ │ ├── Readme.rst │ │ └── docker_build.sh │ ├── end2end │ │ ├── Makefile │ │ ├── end2end.sh │ │ ├── functional-test-config.yaml │ │ └── run-tests.py │ ├── flyte_golang_compile │ │ ├── Readme.rst │ │ ├── flyte_golang_compile.Template │ │ └── update.sh │ ├── github_workflows │ │ ├── Readme.rst │ │ ├── boilerplate_automation.yml │ │ ├── master.yml │ │ ├── pull_request.yml │ │ ├── stale.yml │ │ └── update.sh │ ├── golang_dockerfile │ │ ├── Dockerfile.GoTemplate │ │ ├── Readme.rst │ │ └── update.sh │ ├── golang_support_tools │ │ ├── go.mod │ │ ├── go.sum │ │ └── tools.go │ ├── golang_test_targets │ │ ├── Makefile │ │ ├── Readme.rst │ │ ├── download_tooling.sh │ │ ├── go-gen.sh │ │ └── goimports │ ├── golangci_file │ │ ├── .golangci.yml │ │ ├── Readme.rst │ │ └── update.sh │ ├── precommit │ │ ├── Makefile │ │ ├── hooks │ │ │ ├── pre-push │ │ │ └── prepare-commit-msg │ │ └── update.sh │ └── pull_request_template │ │ ├── Readme.rst │ │ ├── pull_request_template.md │ │ └── update.sh ├── update.cfg └── update.sh ├── cmd ├── entrypoints │ ├── migrate.go │ ├── root.go │ ├── serve.go │ └── serve_dummy.go └── main.go ├── datacatalog.json ├── datacatalog_config.yaml ├── go.mod ├── go.sum ├── pkg ├── common │ └── filters.go ├── config │ ├── config.go │ ├── config_flags.go │ └── config_flags_test.go ├── errors │ ├── errors.go │ └── errors_test.go ├── manager │ ├── impl │ │ ├── artifact_data_store.go │ │ ├── artifact_manager.go │ │ ├── artifact_manager_test.go │ │ ├── dataset_manager.go │ │ ├── dataset_manager_test.go │ │ ├── reservation_manager.go │ │ ├── reservation_manager_test.go │ │ ├── tag_manager.go │ │ ├── tag_manager_test.go │ │ └── validators │ │ │ ├── artifact_validator.go │ │ │ ├── common.go │ │ │ ├── dataset_validator.go │ │ │ ├── errors.go │ │ │ ├── pagination_validator.go │ │ │ ├── partition_validator.go │ │ │ └── tag_validator.go │ ├── interfaces │ │ ├── artifact.go │ │ ├── dataset.go │ │ ├── reservation.go │ │ └── tag.go │ └── mocks │ │ ├── artifact_manager.go │ │ ├── dataset_manager.go │ │ ├── reservation_manager.go │ │ └── tag_manager.go ├── repositories │ ├── config │ │ ├── postgres.go │ │ └── postgres_test.go │ ├── errors │ │ ├── error_transformer.go │ │ ├── errors.go │ │ └── postgres.go │ ├── factory.go │ ├── gormimpl │ │ ├── artifact.go │ │ ├── artifact_test.go │ │ ├── dataset.go │ │ ├── dataset_test.go │ │ ├── filter.go │ │ ├── filter_test.go │ │ ├── join.go │ │ ├── join_test.go │ │ ├── list.go │ │ ├── list_test.go │ │ ├── metrics.go │ │ ├── reservation.go │ │ ├── reservation_test.go │ │ ├── sort.go │ │ ├── sort_test.go │ │ ├── tag.go │ │ └── tag_test.go │ ├── handle.go │ ├── handle_test.go │ ├── initialize.go │ ├── interfaces │ │ ├── artifact_repo.go │ │ ├── base.go │ │ ├── dataset_repo.go │ │ ├── partition_repo.go │ │ ├── reservation_repo.go │ │ └── tag_repo.go │ ├── mocks │ │ ├── artifact_repo.go │ │ ├── base.go │ │ ├── dataset_repo.go │ │ ├── partition_repo.go │ │ ├── reservation_repo.go │ │ └── tag_repo.go │ ├── models │ │ ├── artifact.go │ │ ├── base.go │ │ ├── dataset.go │ │ ├── list.go │ │ ├── partition.go │ │ ├── reservation.go │ │ └── tag.go │ ├── postgres_repo.go │ ├── transformers │ │ ├── artifact.go │ │ ├── artifact_test.go │ │ ├── dataset.go │ │ ├── dataset_test.go │ │ ├── filters.go │ │ ├── filters_test.go │ │ ├── pagination.go │ │ ├── pagination_test.go │ │ ├── partition.go │ │ ├── reservation.go │ │ ├── reservation_test.go │ │ ├── tag.go │ │ ├── tag_test.go │ │ ├── util.go │ │ └── util_test.go │ └── utils │ │ └── test_utils.go ├── rpc │ └── datacatalogservice │ │ └── service.go └── runtime │ ├── application_config_provider.go │ ├── configs │ ├── data_catalog_config.go │ ├── datacatalogconfig_flags.go │ └── datacatalogconfig_flags_test.go │ └── configuration_provider.go └── pull_request_template.md /.github/config.yml: -------------------------------------------------------------------------------- 1 | # Comment to be posted on PRs from first-time contributors in your repository 2 | newPRWelcomeComment: > 3 | Thank you for opening this pull request! 🙌 4 | 5 | These tips will help get your PR across the finish line: 6 | - Most of the repos have a PR template; if not, fill it out to the best of your knowledge. 7 | - Sign off your commits (Reference: [DCO Guide](https://github.com/src-d/guide/blob/master/developer-community/fix-DCO.md)). 8 | 9 | # Comment to be posted to on pull requests merged by a first time user 10 | firstPRMergeComment: > 11 | Congrats on merging your first pull request! 🎉 12 | 13 | # Comment to be posted on first-time issues 14 | newIssueWelcomeComment: > 15 | Thank you for opening your first issue here! 🛠 16 | -------------------------------------------------------------------------------- /.github/workflows/checks.yml: -------------------------------------------------------------------------------- 1 | name: Datacatalog Checks 2 | 3 | on: 4 | workflow_dispatch: 5 | pull_request: 6 | branches: 7 | - master 8 | push: 9 | branches: 10 | - master 11 | 12 | jobs: 13 | lint: 14 | name: Lint 15 | uses: flyteorg/flytetools/.github/workflows/lint.yml@master 16 | with: 17 | go-version: "1.19" 18 | 19 | tests: 20 | name: Unit Tests 21 | uses: flyteorg/flytetools/.github/workflows/tests.yml@master 22 | with: 23 | go-version: "1.19" 24 | secrets: 25 | FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }} 26 | 27 | docker_build: 28 | name: Docker Build Images 29 | uses: flyteorg/flytetools/.github/workflows/docker_build.yml@master 30 | 31 | generate: 32 | name: Check Go Gennerate 33 | uses: flyteorg/flytetools/.github/workflows/go_generate.yml@master 34 | with: 35 | go-version: "1.19" 36 | 37 | bump_version: 38 | name: Bump Version 39 | if: ${{ github.event_name != 'pull_request' }} 40 | needs: [ docker_build, lint, tests, generate ] # Only to ensure it can successfully build 41 | uses: flyteorg/flytetools/.github/workflows/bump_version.yml@master 42 | secrets: 43 | FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }} 44 | 45 | goreleaser: 46 | name: Goreleaser 47 | needs: [ bump_version ] # Only to ensure it can successfully build 48 | uses: flyteorg/flytetools/.github/workflows/goreleaser.yml@master 49 | secrets: 50 | FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }} 51 | with: 52 | go-version: "1.19" 53 | 54 | push_docker_image: 55 | name: Build & Push Datacatalog Image 56 | needs: [ bump_version ] 57 | uses: flyteorg/flytetools/.github/workflows/publish.yml@master 58 | with: 59 | version: ${{ needs.bump_version.outputs.version }} 60 | dockerfile: Dockerfile 61 | push: true 62 | repository: ${{ github.repository }} 63 | secrets: 64 | FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }} 65 | FLYTE_BOT_USERNAME: ${{ secrets.FLYTE_BOT_USERNAME }} 66 | 67 | -------------------------------------------------------------------------------- /.github/workflows/upgrade_automation.yml: -------------------------------------------------------------------------------- 1 | name: Upgrade Automation 2 | on: 3 | workflow_dispatch: 4 | inputs: 5 | component: 6 | description: "Name of Flyte component" 7 | required: true 8 | default: "boilerplate" 9 | type: choice 10 | options: 11 | - boilerplate 12 | - flyteidl 13 | jobs: 14 | trigger-upgrade: 15 | name: ${{ github.event.inputs.component }} Upgrade 16 | uses: flyteorg/flytetools/.github/workflows/flyte_automation.yml@master 17 | with: 18 | component: ${{ github.event.inputs.component }} 19 | go-version: "1.19" 20 | secrets: 21 | FLYTE_BOT_PAT: ${{ secrets.FLYTE_BOT_PAT }} 22 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | vendor/ 3 | vendor-new/ 4 | .DS_Store 5 | bin/ 6 | -------------------------------------------------------------------------------- /.golangci.yml: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | run: 7 | skip-dirs: 8 | - pkg/client 9 | 10 | linters: 11 | disable-all: true 12 | enable: 13 | - deadcode 14 | - errcheck 15 | - gas 16 | - goconst 17 | - goimports 18 | - golint 19 | - gosimple 20 | - govet 21 | - ineffassign 22 | - misspell 23 | - nakedret 24 | - staticcheck 25 | - structcheck 26 | - typecheck 27 | - unconvert 28 | - unparam 29 | - unused 30 | - varcheck 31 | -------------------------------------------------------------------------------- /.goreleaser.yml: -------------------------------------------------------------------------------- 1 | project_name: datacatalog 2 | before: 3 | hooks: 4 | - go mod download 5 | builds: 6 | - id: datacatalog 7 | env: 8 | - CGO_ENABLED=0 9 | main: ./cmd/main.go 10 | binary: datacatalog 11 | goos: 12 | - linux 13 | - windows 14 | - darwin 15 | archives: 16 | - id: datacatalog-archive 17 | name_template: |- 18 | datacatalog_{{ .Tag }}_{{ .Os }}_ 19 | {{- if eq .Arch "amd64" }}x86_64 20 | {{- else if eq .Arch "386" }}i386 21 | {{- else }}{{ .Arch }}{{ end }} 22 | builds: 23 | - datacatalog 24 | format_overrides: 25 | - goos: windows 26 | format: zip 27 | 28 | checksum: 29 | name_template: 'checksums.txt' 30 | changelog: 31 | sort: asc 32 | filters: 33 | exclude: 34 | - '^docs:' 35 | - '^test:' 36 | -------------------------------------------------------------------------------- /CODEOWNERS: -------------------------------------------------------------------------------- 1 | # These are the assigned owners of the datacatalog repo. 2 | # When a PR is opened they will be notified. 3 | * @kumare3 @wild-endeavor 4 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | This project is governed by LF AI Foundation's [code of conduct](https://lfprojects.org/policies/code-of-conduct/). 2 | All contributors and participants agree to abide by its terms. 3 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | FROM --platform=${BUILDPLATFORM} golang:1.19-alpine3.16 as builder 7 | 8 | ARG TARGETARCH 9 | ENV GOARCH "${TARGETARCH}" 10 | ENV GOOS linux 11 | 12 | RUN apk add git openssh-client make curl 13 | 14 | # Create the artifacts directory 15 | RUN mkdir /artifacts 16 | 17 | # Pull GRPC health probe binary for liveness and readiness checks 18 | RUN GRPC_HEALTH_PROBE_VERSION=v0.4.11 && \ 19 | wget -qO/artifacts/grpc_health_probe https://github.com/grpc-ecosystem/grpc-health-probe/releases/download/${GRPC_HEALTH_PROBE_VERSION}/grpc_health_probe-linux-amd64 && \ 20 | chmod +x /artifacts/grpc_health_probe && \ 21 | echo 'ded15e598d887ccc47bf2321371950bbf930f5e4856b9f75712ce4b2b5120480 /artifacts/grpc_health_probe' > .grpc_checksum && \ 22 | sha256sum -c .grpc_checksum 23 | 24 | # COPY only the go mod files for efficient caching 25 | COPY go.mod go.sum /go/src/github.com/flyteorg/datacatalog/ 26 | WORKDIR /go/src/github.com/flyteorg/datacatalog 27 | 28 | # Pull dependencies 29 | RUN go mod download 30 | 31 | # COPY the rest of the source code 32 | COPY . /go/src/github.com/flyteorg/datacatalog/ 33 | 34 | # This 'linux_compile' target should compile binaries to the /artifacts directory 35 | # The main entrypoint should be compiled to /artifacts/datacatalog 36 | RUN make linux_compile 37 | 38 | # update the PATH to include the /artifacts directory 39 | ENV PATH="/artifacts:${PATH}" 40 | 41 | # This will eventually move to centurylink/ca-certs:latest for minimum possible image size 42 | FROM alpine:3.16 43 | LABEL org.opencontainers.image.source=https://github.com/flyteorg/datacatalog 44 | 45 | COPY --from=builder /artifacts /bin 46 | 47 | # Ensure the latest CA certs are present to authenticate SSL connections. 48 | RUN apk --update add ca-certificates 49 | 50 | RUN addgroup -S flyte && adduser -S flyte -G flyte 51 | USER flyte 52 | 53 | CMD ["datacatalog"] 54 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | export REPOSITORY=datacatalog 2 | include boilerplate/flyte/docker_build/Makefile 3 | include boilerplate/flyte/golang_test_targets/Makefile 4 | 5 | .PHONY: update_boilerplate 6 | update_boilerplate: 7 | @curl https://raw.githubusercontent.com/flyteorg/boilerplate/master/boilerplate/update.sh -o boilerplate/update.sh 8 | @boilerplate/update.sh 9 | 10 | .PHONY: compile 11 | compile: 12 | mkdir -p ./bin 13 | go build -o datacatalog ./cmd/main.go && mv ./datacatalog ./bin 14 | 15 | .PHONY: linux_compile 16 | linux_compile: export CGO_ENABLED ?= 0 17 | linux_compile: export GOOS ?= linux 18 | linux_compile: 19 | go build -o /artifacts/datacatalog ./cmd/ 20 | 21 | .PHONY: generate 22 | generate: 23 | @go generate ./... 24 | -------------------------------------------------------------------------------- /NOTICE: -------------------------------------------------------------------------------- 1 | datacatalog 2 | Copyright 2019-2020 Lyft Inc. 3 | 4 | This product includes software developed at Lyft Inc. 5 | 6 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | | 🗑 As of Oct-23 we moved the development of this component to the [monorepo](https://github.com/flyteorg/flyte). 🗑 | 2 | | - | 3 | 4 | # datacatalog 5 | Service that catalogs data to allow for data discovery, lineage and tagging 6 | -------------------------------------------------------------------------------- /boilerplate/flyte/code_of_conduct/CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | This project is governed by LF AI Foundation's [code of conduct](https://lfprojects.org/policies/code-of-conduct/). 2 | All contributors and participants agree to abide by its terms. 3 | -------------------------------------------------------------------------------- /boilerplate/flyte/code_of_conduct/README.rst: -------------------------------------------------------------------------------- 1 | CODE OF CONDUCT 2 | ~~~~~~~~~~~~~~~ 3 | -------------------------------------------------------------------------------- /boilerplate/flyte/code_of_conduct/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | cp ${DIR}/CODE_OF_CONDUCT.md ${DIR}/../../../CODE_OF_CONDUCT.md 13 | -------------------------------------------------------------------------------- /boilerplate/flyte/docker_build/Makefile: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | .PHONY: docker_build 7 | docker_build: 8 | IMAGE_NAME=$$REPOSITORY ./boilerplate/flyte/docker_build/docker_build.sh 9 | 10 | .PHONY: dockerhub_push 11 | dockerhub_push: 12 | IMAGE_NAME=flyteorg/$$REPOSITORY REGISTRY=docker.io ./boilerplate/flyte/docker_build/docker_build.sh 13 | -------------------------------------------------------------------------------- /boilerplate/flyte/docker_build/Readme.rst: -------------------------------------------------------------------------------- 1 | Docker Build and Push 2 | ~~~~~~~~~~~~~~~~~~~~~ 3 | 4 | Provides a ``make docker_build`` target that builds your image locally. 5 | 6 | Provides a ``make dockerhub_push`` target that pushes your final image to Dockerhub. 7 | 8 | The Dockerhub image will tagged ``:`` 9 | 10 | If git head has a git tag, the Dockerhub image will also be tagged ``:``. 11 | 12 | **To Enable:** 13 | 14 | Add ``flyteorg/docker_build`` to your ``boilerplate/update.cfg`` file. 15 | 16 | Add ``include boilerplate/flyte/docker_build/Makefile`` in your main ``Makefile`` _after_ your REPOSITORY environment variable 17 | 18 | :: 19 | 20 | REPOSITORY= 21 | include boilerplate/flyte/docker_build/Makefile 22 | 23 | (this ensures the extra Make targets get included in your main Makefile) 24 | -------------------------------------------------------------------------------- /boilerplate/flyte/docker_build/docker_build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | echo "" 11 | echo "------------------------------------" 12 | echo " DOCKER BUILD" 13 | echo "------------------------------------" 14 | echo "" 15 | 16 | if [ -n "$REGISTRY" ]; then 17 | # Do not push if there are unstaged git changes 18 | CHANGED=$(git status --porcelain) 19 | if [ -n "$CHANGED" ]; then 20 | echo "Please commit git changes before pushing to a registry" 21 | exit 1 22 | fi 23 | fi 24 | 25 | 26 | GIT_SHA=$(git rev-parse HEAD) 27 | 28 | IMAGE_TAG_SUFFIX="" 29 | # for intermediate build phases, append -$BUILD_PHASE to all image tags 30 | if [ -n "$BUILD_PHASE" ]; then 31 | IMAGE_TAG_SUFFIX="-${BUILD_PHASE}" 32 | fi 33 | 34 | IMAGE_TAG_WITH_SHA="${IMAGE_NAME}:${GIT_SHA}${IMAGE_TAG_SUFFIX}" 35 | 36 | RELEASE_SEMVER=$(git describe --tags --exact-match "$GIT_SHA" 2>/dev/null) || true 37 | if [ -n "$RELEASE_SEMVER" ]; then 38 | IMAGE_TAG_WITH_SEMVER="${IMAGE_NAME}:${RELEASE_SEMVER}${IMAGE_TAG_SUFFIX}" 39 | fi 40 | 41 | # build the image 42 | # passing no build phase will build the final image 43 | docker build -t "$IMAGE_TAG_WITH_SHA" --target=${BUILD_PHASE} . 44 | echo "${IMAGE_TAG_WITH_SHA} built locally." 45 | 46 | # if REGISTRY specified, push the images to the remote registy 47 | if [ -n "$REGISTRY" ]; then 48 | 49 | if [ -n "${DOCKER_REGISTRY_PASSWORD}" ]; then 50 | docker login --username="$DOCKER_REGISTRY_USERNAME" --password="$DOCKER_REGISTRY_PASSWORD" 51 | fi 52 | 53 | docker tag "$IMAGE_TAG_WITH_SHA" "${REGISTRY}/${IMAGE_TAG_WITH_SHA}" 54 | 55 | docker push "${REGISTRY}/${IMAGE_TAG_WITH_SHA}" 56 | echo "${REGISTRY}/${IMAGE_TAG_WITH_SHA} pushed to remote." 57 | 58 | # If the current commit has a semver tag, also push the images with the semver tag 59 | if [ -n "$RELEASE_SEMVER" ]; then 60 | 61 | docker tag "$IMAGE_TAG_WITH_SHA" "${REGISTRY}/${IMAGE_TAG_WITH_SEMVER}" 62 | 63 | docker push "${REGISTRY}/${IMAGE_TAG_WITH_SEMVER}" 64 | echo "${REGISTRY}/${IMAGE_TAG_WITH_SEMVER} pushed to remote." 65 | 66 | fi 67 | fi 68 | -------------------------------------------------------------------------------- /boilerplate/flyte/end2end/Makefile: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | .PHONY: end2end_execute 7 | end2end_execute: export FLYTESNACKS_PRIORITIES ?= P0 8 | end2end_execute: export FLYTESNACKS_VERSION ?= $(shell curl --silent "https://api.github.com/repos/flyteorg/flytesnacks/releases/latest" | jq -r .tag_name) 9 | end2end_execute: 10 | ./boilerplate/flyte/end2end/end2end.sh ./boilerplate/flyte/end2end/functional-test-config.yaml --return_non_zero_on_failure 11 | 12 | .PHONY: k8s_integration_execute 13 | k8s_integration_execute: 14 | echo "pass" 15 | -------------------------------------------------------------------------------- /boilerplate/flyte/end2end/end2end.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | set -eu 8 | 9 | CONFIG_FILE=$1; shift 10 | EXTRA_FLAGS=( "$@" ) 11 | 12 | python ./boilerplate/flyte/end2end/run-tests.py $FLYTESNACKS_VERSION $FLYTESNACKS_PRIORITIES $CONFIG_FILE ${EXTRA_FLAGS[@]} 13 | -------------------------------------------------------------------------------- /boilerplate/flyte/end2end/functional-test-config.yaml: -------------------------------------------------------------------------------- 1 | admin: 2 | # For GRPC endpoints you might want to use dns:///flyte.myexample.com 3 | endpoint: dns:///localhost:30080 4 | authType: Pkce 5 | insecure: true 6 | -------------------------------------------------------------------------------- /boilerplate/flyte/flyte_golang_compile/Readme.rst: -------------------------------------------------------------------------------- 1 | Flyte Golang Compile 2 | ~~~~~~~~~~~~~~~~~~~~ 3 | 4 | Common compile script for Flyte golang services. 5 | 6 | **To Enable:** 7 | 8 | Add ``flyteorg/flyte_golang_compile`` to your ``boilerplate/update.cfg`` file. 9 | 10 | Add the following to your Makefile 11 | 12 | :: 13 | 14 | .PHONY: compile_linux 15 | compile_linux: 16 | PACKAGES={{ *your packages }} OUTPUT={{ /path/to/output }} ./boilerplate/flyte/flyte_golang_compile.sh 17 | -------------------------------------------------------------------------------- /boilerplate/flyte/flyte_golang_compile/flyte_golang_compile.Template: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | if [ -z "$PACKAGES" ]; then 7 | echo "PACKAGES environment VAR not set" 8 | exit 1 9 | fi 10 | 11 | if [ -z "$OUTPUT" ]; then 12 | echo "OUTPUT environment VAR not set" 13 | exit 1 14 | fi 15 | 16 | # get the GIT_SHA and RELEASE_SEMVER 17 | 18 | GIT_SHA=$(git rev-parse HEAD) 19 | RELEASE_SEMVER=$(git describe --tags --exact-match $GIT_SHA 2>/dev/null) 20 | 21 | CURRENT_PKG=github.com/flyteorg/{{ REPOSITORY }} 22 | VERSION_PKG="${CURRENT_PKG}/vendor/github.com/flyteorg/flytestdlib" 23 | 24 | LDFLAGS="-X ${VERSION_PKG}/version.Build=${GIT_SHA} -X ${VERSION_PKG}/version.Version=${RELEASE_SEMVER}" 25 | 26 | GOOS=linux GOARCH=amd64 CGO_ENABLED=0 go build -ldflags "$LDFLAGS" -o "$OUTPUT" "$PACKAGES" 27 | -------------------------------------------------------------------------------- /boilerplate/flyte/flyte_golang_compile/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | echo " - generating ${DIR}/flyte_golang_compile.sh" 13 | sed -e "s/{{REPOSITORY}}/${REPOSITORY}/g" ${DIR}/flyte_golang_compile.Template > ${DIR}/flyte_golang_compile.sh 14 | -------------------------------------------------------------------------------- /boilerplate/flyte/github_workflows/Readme.rst: -------------------------------------------------------------------------------- 1 | Golang Github Actions 2 | ~~~~~~~~~~~~~~~~~ 3 | 4 | Provides a two github actions workflows. 5 | 6 | **To Enable:** 7 | 8 | Add ``flyteorg/github_workflows`` to your ``boilerplate/update.cfg`` file. 9 | 10 | Add a github secret ``package_name`` with the name to use for publishing (e.g. ``flytepropeller``). Typicaly, this will be the same name as the repository. 11 | 12 | *Note*: If you are working on a fork, include that prefix in your package name (``myfork/flytepropeller``). 13 | 14 | The actions will push to 2 repos: 15 | 16 | 1. ``docker.pkg.github.com/flyteorg//`` 17 | 2. ``docker.pkg.github.com/flyteorg//-stages`` : this repo is used to cache build stages to speed up iterative builds after. 18 | 19 | There are two workflows that get deployed: 20 | 21 | 1. A workflow that runs on Pull Requests to build and push images to github registy tagged with the commit sha. 22 | 2. A workflow that runs on master merges that bump the patch version of release tag, builds and pushes images to github registry tagged with the version, commit sha as well as "latest" 23 | -------------------------------------------------------------------------------- /boilerplate/flyte/github_workflows/boilerplate_automation.yml: -------------------------------------------------------------------------------- 1 | name: Update Boilerplate Automation 2 | on: 3 | workflow_dispatch: 4 | jobs: 5 | update-boilerplate: 6 | name: Update Boilerplate 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/checkout@v2 10 | with: 11 | fetch-depth: "0" 12 | - name: Update Boilerplate 13 | run: | 14 | make update_boilerplate 15 | - name: Create Pull Request 16 | id: cpr 17 | uses: peter-evans/create-pull-request@v3 18 | with: 19 | token: ${{ secrets.FLYTE_BOT_PAT }} 20 | commit-message: Update Boilerplate 21 | committer: Flyte-Bot 22 | author: ${{ github.actor }} <${{ github.actor }}@users.noreply.github.com> 23 | signoff: true 24 | branch: flyte-bot-update-boilerplate 25 | delete-branch: true 26 | title: 'Update Boilerplate' 27 | body: | 28 | Update Boilerplate 29 | - Auto-generated by [flyte-bot] 30 | labels: | 31 | boilerplate 32 | team-reviewers: | 33 | owners 34 | maintainers 35 | draft: false 36 | 37 | -------------------------------------------------------------------------------- /boilerplate/flyte/github_workflows/master.yml: -------------------------------------------------------------------------------- 1 | name: Master 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | 8 | jobs: 9 | build: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@master 13 | with: 14 | fetch-depth: '0' 15 | - name: Bump version and push tag 16 | id: bump-version 17 | uses: anothrNick/github-tag-action@1.17.2 18 | env: 19 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 20 | WITH_V: true 21 | DEFAULT_BUMP: patch 22 | - name: Push Docker Image to Github Registry 23 | uses: whoan/docker-build-with-cache-action@v5 24 | with: 25 | username: "${{ github.actor }}" 26 | password: "${{ secrets.GITHUB_TOKEN }}" 27 | image_name: ${{ secrets.package_name }} 28 | image_tag: latest,${{ github.sha }},${{ steps.bump-version.outputs.tag }} 29 | push_git_tag: true 30 | registry: docker.pkg.github.com 31 | build_extra_args: "--compress=true" 32 | -------------------------------------------------------------------------------- /boilerplate/flyte/github_workflows/pull_request.yml: -------------------------------------------------------------------------------- 1 | name: Pull Request 2 | 3 | on: 4 | pull_request 5 | 6 | jobs: 7 | build: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v1 11 | - name: Push Docker Image to Github Registry 12 | uses: whoan/docker-build-with-cache-action@v5 13 | with: 14 | username: "${{ github.actor }}" 15 | password: "${{ secrets.GITHUB_TOKEN }}" 16 | image_name: ${{ secrets.package_name }} 17 | image_tag: ${{ github.sha }} 18 | push_git_tag: true 19 | registry: docker.pkg.github.com 20 | -------------------------------------------------------------------------------- /boilerplate/flyte/github_workflows/stale.yml: -------------------------------------------------------------------------------- 1 | # Configuration for probot-stale - https://github.com/probot/stale 2 | 3 | # Number of days of inactivity before an Issue or Pull Request becomes stale 4 | daysUntilStale: 120 5 | 6 | # Number of days of inactivity before an Issue or Pull Request with the stale label is closed. 7 | # Set to false to disable. If disabled, issues still need to be closed manually, but will remain marked as stale. 8 | daysUntilClose: 7 9 | 10 | # Only issues or pull requests with all of these labels are check if stale. Defaults to `[]` (disabled) 11 | onlyLabels: [] 12 | 13 | # Issues or Pull Requests with these labels will never be considered stale. Set to `[]` to disable 14 | exemptLabels: 15 | - pinned 16 | - security 17 | - "[Status] Maybe Later" 18 | 19 | # Set to true to ignore issues in a project (defaults to false) 20 | exemptProjects: false 21 | 22 | # Set to true to ignore issues in a milestone (defaults to false) 23 | exemptMilestones: false 24 | 25 | # Set to true to ignore issues with an assignee (defaults to false) 26 | exemptAssignees: false 27 | 28 | # Label to use when marking as stale 29 | staleLabel: wontfix 30 | 31 | # Comment to post when marking as stale. Set to `false` to disable 32 | markComment: > 33 | This issue/pullrequest has been automatically marked as stale because it has not had 34 | recent activity. It will be closed if no further activity occurs. Thank you 35 | for your contributions. 36 | 37 | # Comment to post when removing the stale label. 38 | # unmarkComment: > 39 | # Your comment here. 40 | 41 | # Comment to post when closing a stale Issue or Pull Request. 42 | # closeComment: > 43 | # Your comment here. 44 | 45 | # Limit the number of actions per hour, from 1-30. Default is 30 46 | limitPerRun: 30 47 | 48 | # Limit to only `issues` or `pulls` 49 | only: pulls 50 | 51 | # Optionally, specify configuration settings that are specific to just 'issues' or 'pulls': 52 | # pulls: 53 | # daysUntilStale: 30 54 | # markComment: > 55 | # This pull request has been automatically marked as stale because it has not had 56 | # recent activity. It will be closed if no further activity occurs. Thank you 57 | # for your contributions. 58 | 59 | # issues: 60 | # exemptLabels: 61 | # - confirmed -------------------------------------------------------------------------------- /boilerplate/flyte/github_workflows/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | mkdir -p ${DIR}/../../../.github/workflows 13 | 14 | echo " - generating github action workflows in root directory." 15 | sed -e "s/{{REPOSITORY}}/${REPOSITORY}/g" ${DIR}/master.yml > ${DIR}/../../../.github/workflows/master.yml 16 | sed -e "s/{{REPOSITORY}}/${REPOSITORY}/g" ${DIR}/pull_request.yml > ${DIR}/../../../.github/workflows/pull_request.yml 17 | cp ${DIR}/stale.yml ${DIR}/../../../.github/stale.yml 18 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_dockerfile/Dockerfile.GoTemplate: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | FROM golang:1.17.1-alpine3.14 as builder 7 | RUN apk add git openssh-client make curl 8 | 9 | # COPY only the go mod files for efficient caching 10 | COPY go.mod go.sum /go/src/github.com/flyteorg/{{REPOSITORY}}/ 11 | WORKDIR /go/src/github.com/flyteorg/{{REPOSITORY}} 12 | 13 | # Pull dependencies 14 | RUN go mod download 15 | 16 | # COPY the rest of the source code 17 | COPY . /go/src/github.com/flyteorg/{{REPOSITORY}}/ 18 | 19 | # This 'linux_compile' target should compile binaries to the /artifacts directory 20 | # The main entrypoint should be compiled to /artifacts/{{REPOSITORY}} 21 | RUN make linux_compile 22 | 23 | # install grpc-health-probe 24 | RUN curl --silent --fail --show-error --location --output /artifacts/grpc_health_probe "https://github.com/grpc-ecosystem/grpc-health-probe/releases/download/v0.4.5/grpc_health_probe-linux-amd64" && \ 25 | chmod +x /artifacts/grpc_health_probe && \ 26 | echo '8699c46352d752d8f533cae72728b0e65663f399fc28fb9cd854b14ad5f85f44 /artifacts/grpc_health_probe' > .grpc_checksum && \ 27 | sha256sum -c .grpc_checksum 28 | 29 | # update the PATH to include the /artifacts directory 30 | ENV PATH="/artifacts:${PATH}" 31 | 32 | # This will eventually move to centurylink/ca-certs:latest for minimum possible image size 33 | FROM alpine:3.14 34 | COPY --from=builder /artifacts /bin 35 | 36 | RUN apk --update add ca-certificates 37 | 38 | CMD ["{{REPOSITORY}}"] 39 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_dockerfile/Readme.rst: -------------------------------------------------------------------------------- 1 | Golang Dockerfile 2 | ~~~~~~~~~~~~~~~~~ 3 | 4 | Provides a Dockerfile that produces a small image. 5 | 6 | **To Enable:** 7 | 8 | Add ``flyteorg/golang_dockerfile`` to your ``boilerplate/update.cfg`` file. 9 | 10 | Create and configure a ``make linux_compile`` target that compiles your go binaries to the ``/artifacts`` directory :: 11 | 12 | .PHONY: linux_compile 13 | linux_compile: 14 | RUN GOOS=linux GOARCH=amd64 CGO_ENABLED=0 go build -o /artifacts {{ packages }} 15 | 16 | All binaries compiled to ``/artifacts`` will be available at ``/bin`` in your final image. 17 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_dockerfile/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | echo " - generating Dockerfile in root directory." 13 | sed -e "s/{{REPOSITORY}}/${REPOSITORY}/g" ${DIR}/Dockerfile.GoTemplate > ${DIR}/../../../Dockerfile 14 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_support_tools/tools.go: -------------------------------------------------------------------------------- 1 | //go:build tools 2 | // +build tools 3 | 4 | package tools 5 | 6 | import ( 7 | _ "github.com/EngHabu/mockery/cmd/mockery" 8 | _ "github.com/alvaroloes/enumer" 9 | _ "github.com/flyteorg/flytestdlib/cli/pflags" 10 | _ "github.com/golangci/golangci-lint/cmd/golangci-lint" 11 | _ "github.com/pseudomuto/protoc-gen-doc/cmd/protoc-gen-doc" 12 | ) 13 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_test_targets/Makefile: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | 7 | .PHONY: download_tooling 8 | download_tooling: #download dependencies (including test deps) for the package 9 | @boilerplate/flyte/golang_test_targets/download_tooling.sh 10 | 11 | .PHONY: generate 12 | generate: download_tooling #generate go code 13 | @boilerplate/flyte/golang_test_targets/go-gen.sh 14 | 15 | .PHONY: lint 16 | lint: download_tooling #lints the package for common code smells 17 | GL_DEBUG=linters_output,env golangci-lint run --deadline=5m --exclude deprecated -v 18 | 19 | # If code is failing goimports linter, this will fix. 20 | # skips 'vendor' 21 | .PHONY: goimports 22 | goimports: 23 | @boilerplate/flyte/golang_test_targets/goimports 24 | 25 | .PHONY: mod_download 26 | mod_download: #download dependencies (including test deps) for the package 27 | go mod download 28 | 29 | .PHONY: install 30 | install: download_tooling mod_download 31 | 32 | .PHONY: show 33 | show: 34 | go list -m all 35 | 36 | .PHONY: test_unit 37 | test_unit: 38 | go test -cover ./... -race 39 | 40 | .PHONY: test_benchmark 41 | test_benchmark: 42 | go test -bench . ./... 43 | 44 | .PHONY: test_unit_cover 45 | test_unit_cover: 46 | go test ./... -coverprofile /tmp/cover.out -covermode=count 47 | go tool cover -func /tmp/cover.out 48 | 49 | .PHONY: test_unit_visual 50 | test_unit_visual: 51 | go test ./... -coverprofile /tmp/cover.out -covermode=count 52 | go tool cover -html=/tmp/cover.out 53 | 54 | .PHONY: test_unit_codecov 55 | test_unit_codecov: 56 | go test ./... -race -coverprofile=coverage.txt -covermode=atomic 57 | curl -s https://codecov.io/bash > codecov_bash.sh && bash codecov_bash.sh 58 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_test_targets/Readme.rst: -------------------------------------------------------------------------------- 1 | Golang Test Targets 2 | ~~~~~~~~~~~~~~~~~~~ 3 | 4 | Provides an ``install`` make target that uses ``go mod`` to install golang dependencies. 5 | 6 | Provides a ``lint`` make target that uses golangci to lint your code. 7 | 8 | Provides a ``test_unit`` target for unit tests. 9 | 10 | Provides a ``test_unit_cover`` target for analysing coverage of unit tests, which will output the coverage of each function and total statement coverage. 11 | 12 | Provides a ``test_unit_visual`` target for visualizing coverage of unit tests through an interactive html code heat map. 13 | 14 | Provides a ``test_benchmark`` target for benchmark tests. 15 | 16 | **To Enable:** 17 | 18 | Add ``flyteorg/golang_test_targets`` to your ``boilerplate/update.cfg`` file. 19 | 20 | Make sure you're using ``go mod`` for dependency management. 21 | 22 | Provide a ``.golangci`` configuration (the lint target requires it). 23 | 24 | Add ``include boilerplate/flyte/golang_test_targets/Makefile`` in your main ``Makefile`` _after_ your REPOSITORY environment variable 25 | 26 | :: 27 | 28 | REPOSITORY= 29 | include boilerplate/flyte/golang_test_targets/Makefile 30 | 31 | (this ensures the extra make targets get included in your main Makefile) 32 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_test_targets/download_tooling.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Everything in this file needs to be installed outside of current module 4 | # The reason we cannot turn off module entirely and install is that we need the replace statement in go.mod 5 | # because we are installing a mockery fork. Turning it off would result installing the original not the fork. 6 | # We also want to version all the other tools. We also want to be able to run go mod tidy without removing the version 7 | # pins. To facilitate this, we're maintaining two sets of go.mod/sum files - the second one only for tooling. This is 8 | # the same approach that go 1.14 will take as well. 9 | # See: 10 | # https://github.com/flyteorg/flyte/issues/129 11 | # https://github.com/golang/go/issues/30515 for some background context 12 | # https://github.com/go-modules-by-example/index/blob/5ec250b4b78114a55001bd7c9cb88f6e07270ea5/010_tools/README.md 13 | 14 | set -e 15 | 16 | # List of tools to go get 17 | # In the format of ":" or ":" if no cli 18 | tools=( 19 | "github.com/EngHabu/mockery/cmd/mockery" 20 | "github.com/flyteorg/flytestdlib/cli/pflags@latest" 21 | "github.com/golangci/golangci-lint/cmd/golangci-lint" 22 | "github.com/alvaroloes/enumer" 23 | "github.com/pseudomuto/protoc-gen-doc/cmd/protoc-gen-doc" 24 | ) 25 | 26 | tmp_dir=$(mktemp -d -t gotooling-XXX) 27 | echo "Using temp directory ${tmp_dir}" 28 | cp -R boilerplate/flyte/golang_support_tools/* $tmp_dir 29 | pushd "$tmp_dir" 30 | 31 | for tool in "${tools[@]}" 32 | do 33 | echo "Installing ${tool}" 34 | GO111MODULE=on go install $tool 35 | done 36 | 37 | popd 38 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_test_targets/go-gen.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -ex 4 | 5 | echo "Running go generate" 6 | go generate ./... 7 | 8 | # This section is used by GitHub workflow to ensure that the generation step was run 9 | if [ -n "$DELTA_CHECK" ]; then 10 | DIRTY=$(git status --porcelain) 11 | if [ -n "$DIRTY" ]; then 12 | echo "FAILED: Go code updated without commiting generated code." 13 | echo "Ensure make generate has run and all changes are committed." 14 | DIFF=$(git diff) 15 | echo "diff detected: $DIFF" 16 | DIFF=$(git diff --name-only) 17 | echo "files different: $DIFF" 18 | exit 1 19 | else 20 | echo "SUCCESS: Generated code is up to date." 21 | fi 22 | fi 23 | -------------------------------------------------------------------------------- /boilerplate/flyte/golang_test_targets/goimports: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | goimports -w $(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pkg/client/*" -not -path "./boilerplate/*") 9 | -------------------------------------------------------------------------------- /boilerplate/flyte/golangci_file/.golangci.yml: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | run: 7 | skip-dirs: 8 | - pkg/client 9 | 10 | linters: 11 | disable-all: true 12 | enable: 13 | - deadcode 14 | - errcheck 15 | - gas 16 | - goconst 17 | - goimports 18 | - golint 19 | - gosimple 20 | - govet 21 | - ineffassign 22 | - misspell 23 | - nakedret 24 | - staticcheck 25 | - structcheck 26 | - typecheck 27 | - unconvert 28 | - unparam 29 | - unused 30 | - varcheck 31 | -------------------------------------------------------------------------------- /boilerplate/flyte/golangci_file/Readme.rst: -------------------------------------------------------------------------------- 1 | GolangCI File 2 | ~~~~~~~~~~~~~ 3 | 4 | Provides a ``.golangci`` file with the linters we've agreed upon. 5 | 6 | **To Enable:** 7 | 8 | Add ``flyteorg/golangci_file`` to your ``boilerplate/update.cfg`` file. 9 | -------------------------------------------------------------------------------- /boilerplate/flyte/golangci_file/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | # Clone the .golangci file 13 | echo " - copying ${DIR}/.golangci to the root directory." 14 | cp ${DIR}/.golangci.yml ${DIR}/../../../.golangci.yml 15 | -------------------------------------------------------------------------------- /boilerplate/flyte/precommit/Makefile: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | 6 | 7 | .PHONY: setup-precommit 8 | setup-precommit: #setup the precommit 9 | @boilerplate/flyte/precommit/update.sh 10 | -------------------------------------------------------------------------------- /boilerplate/flyte/precommit/hooks/pre-push: -------------------------------------------------------------------------------- 1 | DUMMY_SHA=0000000000000000000000000000000000000000 2 | 3 | echo "Running pre-push check; to skip this step use 'push --no-verify'" 4 | 5 | while read LOCAL_REF LOCAL_SHA REMOTE_REF REMOTE_SHA 6 | do 7 | if [ "$LOCAL_SHA" = $DUMMY_SHA ] 8 | then 9 | # Branch deleted. Do nothing. 10 | exit 0 11 | else 12 | if [ "$REMOTE_SHA" = $DUMMY_SHA ] 13 | then 14 | # New branch. Verify the last commit, since this is very likely where the new code is 15 | # (though there is no way to know for sure). In the extremely uncommon case in which someone 16 | # pushes more than 1 new commit to a branch, CI will enforce full checking. 17 | RANGE="$LOCAL_SHA~1..$LOCAL_SHA" 18 | else 19 | # Updating branch. Verify new commits. 20 | RANGE="$REMOTE_SHA..$LOCAL_SHA" 21 | fi 22 | 23 | # Verify DCO signoff. We do this before the format checker, since it has 24 | # some probability of failing spuriously, while this check never should. 25 | # 26 | # In general, we can't assume that the commits are signed off by author 27 | # pushing, so we settle for just checking that there is a signoff at all. 28 | SIGNED_OFF=$(git rev-list --no-merges --grep "^Signed-off-by: " "$RANGE") 29 | NOT_SIGNED_OFF=$(git rev-list --no-merges "$RANGE" | grep -Fxv "$SIGNED_OFF") 30 | if [ -n "$NOT_SIGNED_OFF" ] 31 | then 32 | echo >&2 "ERROR: The following commits do not have DCO signoff:" 33 | while read -r commit; do 34 | echo " $(git log --pretty=oneline --abbrev-commit -n 1 $commit)" 35 | done <<< "$NOT_SIGNED_OFF" 36 | exit 1 37 | fi 38 | fi 39 | done 40 | 41 | exit 0 42 | -------------------------------------------------------------------------------- /boilerplate/flyte/precommit/hooks/prepare-commit-msg: -------------------------------------------------------------------------------- 1 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 2 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 3 | # 4 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 5 | # $ ln -s ../../support/hooks/prepare-commit-msg .git/hooks/prepare-commit-msg 6 | 7 | COMMIT_MESSAGE_FILE="$1" 8 | AUTHOR=$(git var GIT_AUTHOR_IDENT) 9 | SIGNOFF=$(echo $AUTHOR | sed -n 's/^\(.*>\).*$/Signed-off-by: \1/p') 10 | 11 | # Check for DCO signoff message. If one doesn't exist, append one and then warn 12 | # the user that you did so. 13 | if ! $(grep -qs "^$SIGNOFF" "$COMMIT_MESSAGE_FILE") ; then 14 | echo "\n$SIGNOFF" >> "$COMMIT_MESSAGE_FILE" 15 | echo "Appended the following signoff to the end of the commit message:\n $SIGNOFF\n" 16 | fi 17 | -------------------------------------------------------------------------------- /boilerplate/flyte/precommit/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | # Helper script for Automatically add DCO signoff with commit hooks 11 | # Taken from Envoy https://gitlab.cncf.ci/envoyproxy/envoy 12 | if [ ! "$PWD" == "$(git rev-parse --show-toplevel)" ]; then 13 | cat >&2 <<__EOF__ 14 | ERROR: this script must be run at the root of the envoy source tree 15 | __EOF__ 16 | exit 1 17 | fi 18 | 19 | # Helper functions that calculate `abspath` and `relpath`. Taken from Mesos 20 | # commit 82b040a60561cf94dec3197ea88ae15e57bcaa97, which also carries the Apache 21 | # V2 license, and has deployed this code successfully for some time. 22 | abspath() { 23 | cd "$(dirname "${1}")" 24 | echo "${PWD}"/"$(basename "${1}")" 25 | cd "${OLDPWD}" 26 | } 27 | relpath() { 28 | local FROM TO UP 29 | FROM="$(abspath "${1%/}")" TO="$(abspath "${2%/}"/)" 30 | while test "${TO}" = "${TO#"${FROM}"/}" \ 31 | -a "${TO}" != "${FROM}"; do 32 | FROM="${FROM%/*}" UP="../${UP}" 33 | done 34 | TO="${UP%/}${TO#${FROM}}" 35 | echo "${TO:-.}" 36 | } 37 | 38 | # Try to find the `.git` directory, even if it's not in Flyte project root (as 39 | # it wouldn't be if, say, this were in a submodule). The "blessed" but fairly 40 | # new way to do this is to use `--git-common-dir`. 41 | DOT_GIT_DIR=$(git rev-parse --git-common-dir) 42 | if test ! -d "${DOT_GIT_DIR}"; then 43 | # If `--git-common-dir` is not available, fall back to older way of doing it. 44 | DOT_GIT_DIR=$(git rev-parse --git-dir) 45 | fi 46 | 47 | mkdir -p ${DOT_GIT_DIR}/hooks 48 | 49 | HOOKS_DIR="${DOT_GIT_DIR}/hooks" 50 | HOOKS_DIR_RELPATH=$(relpath "${HOOKS_DIR}" "${PWD}") 51 | 52 | if [ ! -e "${HOOKS_DIR}/prepare-commit-msg" ]; then 53 | echo "Installing hook 'prepare-commit-msg'" 54 | ln -s "${HOOKS_DIR_RELPATH}/boilerplate/flyte/precommit/hooks/prepare-commit-msg" "${HOOKS_DIR}/prepare-commit-msg" 55 | fi 56 | 57 | if [ ! -e "${HOOKS_DIR}/pre-push" ]; then 58 | echo "Installing hook 'pre-push'" 59 | ln -s "${HOOKS_DIR_RELPATH}/boilerplate/flyte/precommit/hooks/pre-push" "${HOOKS_DIR}/pre-push" 60 | fi 61 | -------------------------------------------------------------------------------- /boilerplate/flyte/pull_request_template/Readme.rst: -------------------------------------------------------------------------------- 1 | Pull Request Template 2 | ~~~~~~~~~~~~~~~~~~~~~ 3 | 4 | Provides a Pull Request template. 5 | 6 | **To Enable:** 7 | 8 | Add ``flyteorg/golang_test_targets`` to your ``boilerplate/update.cfg`` file. 9 | -------------------------------------------------------------------------------- /boilerplate/flyte/pull_request_template/pull_request_template.md: -------------------------------------------------------------------------------- 1 | ## _Read then delete this section_ 2 | 3 | _- Make sure to use a concise title for the pull-request._ 4 | 5 | _- Use #patch, #minor or #major in the pull-request title to bump the corresponding version. Otherwise, the patch version 6 | will be bumped. [More details](https://github.com/marketplace/actions/github-tag-bump)_ 7 | 8 | # TL;DR 9 | _Please replace this text with a description of what this PR accomplishes._ 10 | 11 | ## Type 12 | - [ ] Bug Fix 13 | - [ ] Feature 14 | - [ ] Plugin 15 | 16 | ## Are all requirements met? 17 | 18 | - [ ] Code completed 19 | - [ ] Smoke tested 20 | - [ ] Unit tests added 21 | - [ ] Code documentation added 22 | - [ ] Any pending items have an associated Issue 23 | 24 | ## Complete description 25 | _How did you fix the bug, make the feature etc. Link to any design docs etc_ 26 | 27 | ## Tracking Issue 28 | _Remove the '*fixes*' keyword if there will be multiple PRs to fix the linked issue_ 29 | 30 | fixes https://github.com/flyteorg/flyte/issues/ 31 | 32 | ## Follow-up issue 33 | _NA_ 34 | OR 35 | _https://github.com/flyteorg/flyte/issues/_ 36 | -------------------------------------------------------------------------------- /boilerplate/flyte/pull_request_template/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | cp ${DIR}/pull_request_template.md ${DIR}/../../../pull_request_template.md 13 | -------------------------------------------------------------------------------- /boilerplate/update.cfg: -------------------------------------------------------------------------------- 1 | flyte/docker_build 2 | flyte/golang_test_targets 3 | flyte/golangci_file 4 | flyte/golang_support_tools 5 | flyte/pull_request_template 6 | flyte/ 7 | -------------------------------------------------------------------------------- /boilerplate/update.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # WARNING: THIS FILE IS MANAGED IN THE 'BOILERPLATE' REPO AND COPIED TO OTHER REPOSITORIES. 4 | # ONLY EDIT THIS FILE FROM WITHIN THE 'FLYTEORG/BOILERPLATE' REPOSITORY: 5 | # 6 | # TO OPT OUT OF UPDATES, SEE https://github.com/flyteorg/boilerplate/blob/master/Readme.rst 7 | 8 | set -e 9 | 10 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 11 | 12 | OUT="$(mktemp -d)" 13 | trap 'rm -fr $OUT' EXIT 14 | 15 | git clone https://github.com/flyteorg/boilerplate.git "${OUT}" 16 | 17 | echo "Updating the update.sh script." 18 | cp "${OUT}/boilerplate/update.sh" "${DIR}/update.sh" 19 | 20 | CONFIG_FILE="${DIR}/update.cfg" 21 | README="https://github.com/flyteorg/boilerplate/blob/master/Readme.rst" 22 | 23 | if [ ! -f "$CONFIG_FILE" ]; then 24 | echo "$CONFIG_FILE not found." 25 | echo "This file is required in order to select which features to include." 26 | echo "See $README for more details." 27 | exit 1 28 | fi 29 | 30 | if [ -z "$REPOSITORY" ]; then 31 | echo "$REPOSITORY is required to run this script" 32 | echo "See $README for more details." 33 | exit 1 34 | fi 35 | 36 | while read -r directory junk; do 37 | # Skip comment lines (which can have leading whitespace) 38 | if [[ "$directory" == '#'* ]]; then 39 | continue 40 | fi 41 | # Skip blank or whitespace-only lines 42 | if [[ "$directory" == "" ]]; then 43 | continue 44 | fi 45 | # Lines like 46 | # valid/path other_junk 47 | # are not acceptable, unless `other_junk` is a comment 48 | if [[ "$junk" != "" ]] && [[ "$junk" != '#'* ]]; then 49 | echo "Invalid config! Only one directory is allowed per line. Found '$junk'" 50 | exit 1 51 | fi 52 | 53 | dir_path="${OUT}/boilerplate/${directory}" 54 | # Make sure the directory exists 55 | if ! [[ -d "$dir_path" ]]; then 56 | echo "Invalid boilerplate directory: '$directory'" 57 | exit 1 58 | fi 59 | 60 | echo "***********************************************************************************" 61 | echo "$directory is configured in update.cfg." 62 | echo "-----------------------------------------------------------------------------------" 63 | echo "syncing files from source." 64 | rm -rf "${DIR:?}/${directory}" 65 | mkdir -p "$(dirname "${DIR}"/"${directory}")" 66 | cp -r "$dir_path" "${DIR}/${directory}" 67 | if [ -f "${DIR}/${directory}/update.sh" ]; then 68 | echo "executing ${DIR}/${directory}/update.sh" 69 | "${DIR}/${directory}/update.sh" 70 | fi 71 | echo "***********************************************************************************" 72 | echo "" 73 | done < "$CONFIG_FILE" 74 | -------------------------------------------------------------------------------- /cmd/entrypoints/migrate.go: -------------------------------------------------------------------------------- 1 | package entrypoints 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/repositories" 5 | 6 | "context" 7 | 8 | "github.com/spf13/cobra" 9 | ) 10 | 11 | var parentMigrateCmd = &cobra.Command{ 12 | Use: "migrate", 13 | Short: "This command controls migration behavior for the Flyte Catalog database. Please choose a subcommand.", 14 | } 15 | 16 | // This runs all the migrations 17 | var migrateCmd = &cobra.Command{ 18 | Use: "run", 19 | Short: "This command will run all the migrations for the database", 20 | RunE: func(cmd *cobra.Command, args []string) error { 21 | ctx := context.Background() 22 | return repositories.Migrate(ctx) 23 | }, 24 | } 25 | 26 | func init() { 27 | RootCmd.AddCommand(parentMigrateCmd) 28 | parentMigrateCmd.AddCommand(migrateCmd) 29 | } 30 | -------------------------------------------------------------------------------- /cmd/entrypoints/root.go: -------------------------------------------------------------------------------- 1 | package entrypoints 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "fmt" 7 | "os" 8 | 9 | "github.com/flyteorg/flytestdlib/logger" 10 | 11 | "github.com/flyteorg/flytestdlib/config" 12 | "github.com/flyteorg/flytestdlib/config/viper" 13 | "github.com/spf13/cobra" 14 | "github.com/spf13/pflag" 15 | ) 16 | 17 | var ( 18 | cfgFile string 19 | 20 | configAccessor = viper.NewAccessor(config.Options{StrictMode: true}) 21 | ) 22 | 23 | func init() { 24 | // See https://gist.github.com/nak3/78a32817a8a3950ae48f239a44cd3663 25 | // allows `$ datacatalog --logtostderr` to work 26 | pflag.CommandLine.AddGoFlagSet(flag.CommandLine) 27 | 28 | // Add persistent flags - persistent flags persist through all sub-commands 29 | RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is ./datacatalog_config.yaml)") 30 | 31 | RootCmd.AddCommand(viper.GetConfigCommand()) 32 | 33 | // Allow viper to read the value of the flags 34 | configAccessor.InitializePflags(RootCmd.PersistentFlags()) 35 | 36 | err := flag.CommandLine.Parse([]string{}) 37 | if err != nil { 38 | fmt.Println(err) 39 | os.Exit(-1) 40 | } 41 | } 42 | 43 | // Execute adds all child commands to the root command sets flags appropriately. 44 | // This is called by main.main(). It only needs to happen once to the rootCmd. 45 | func Execute() error { 46 | if err := RootCmd.Execute(); err != nil { 47 | fmt.Println(err) 48 | os.Exit(1) 49 | } 50 | return nil 51 | } 52 | 53 | // RootCmd represents the base command when called without any subcommands 54 | var RootCmd = &cobra.Command{ 55 | Use: "datacatalog", 56 | Short: "Launches datacatalog", 57 | Long: ` 58 | To get started run the serve subcommand which will start a server on localhost:8089: 59 | 60 | datacatalog serve 61 | `, 62 | PersistentPreRunE: func(cmd *cobra.Command, args []string) error { 63 | return initConfig(cmd.Flags()) 64 | }, 65 | } 66 | 67 | func initConfig(flags *pflag.FlagSet) error { 68 | configAccessor = viper.NewAccessor(config.Options{ 69 | SearchPaths: []string{cfgFile, ".", "/etc/flyte/config", "$GOPATH/src/github.com/flyteorg/datacatalog"}, 70 | StrictMode: false, 71 | }) 72 | 73 | logger.Infof(context.TODO(), "Using config file: %v", configAccessor.ConfigFilesUsed()) 74 | 75 | configAccessor.InitializePflags(flags) 76 | 77 | return configAccessor.UpdateConfig(context.TODO()) 78 | } 79 | -------------------------------------------------------------------------------- /cmd/entrypoints/serve.go: -------------------------------------------------------------------------------- 1 | package entrypoints 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/config" 7 | "github.com/flyteorg/datacatalog/pkg/rpc/datacatalogservice" 8 | "github.com/flyteorg/datacatalog/pkg/runtime" 9 | "github.com/flyteorg/flytestdlib/contextutils" 10 | "github.com/flyteorg/flytestdlib/logger" 11 | "github.com/flyteorg/flytestdlib/profutils" 12 | "github.com/flyteorg/flytestdlib/promutils/labeled" 13 | 14 | "github.com/spf13/cobra" 15 | ) 16 | 17 | var serveCmd = &cobra.Command{ 18 | Use: "serve", 19 | Short: "Launches the Data Catalog server", 20 | RunE: func(cmd *cobra.Command, args []string) error { 21 | ctx := context.Background() 22 | cfg := config.GetConfig() 23 | 24 | // serve a http healthcheck endpoint 25 | go func() { 26 | err := datacatalogservice.ServeHTTPHealthCheck(ctx, cfg) 27 | if err != nil { 28 | logger.Errorf(ctx, "Unable to serve http", config.GetConfig().GetHTTPHostAddress(), err) 29 | } 30 | }() 31 | 32 | // Serve profiling endpoint. 33 | dataCatalogConfig := runtime.NewConfigurationProvider().ApplicationConfiguration().GetDataCatalogConfig() 34 | go func() { 35 | err := profutils.StartProfilingServerWithDefaultHandlers( 36 | context.Background(), dataCatalogConfig.ProfilerPort, nil) 37 | if err != nil { 38 | logger.Panicf(context.Background(), "Failed to Start profiling and Metrics server. Error, %v", err) 39 | } 40 | }() 41 | 42 | // Set Keys 43 | labeled.SetMetricKeys(contextutils.AppNameKey, contextutils.ProjectKey, contextutils.DomainKey) 44 | 45 | return datacatalogservice.ServeInsecure(ctx, cfg) 46 | }, 47 | } 48 | 49 | func init() { 50 | RootCmd.AddCommand(serveCmd) 51 | } 52 | -------------------------------------------------------------------------------- /cmd/entrypoints/serve_dummy.go: -------------------------------------------------------------------------------- 1 | package entrypoints 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/config" 7 | "github.com/flyteorg/datacatalog/pkg/rpc/datacatalogservice" 8 | "github.com/flyteorg/flytestdlib/logger" 9 | "github.com/spf13/cobra" 10 | ) 11 | 12 | var serveDummyCmd = &cobra.Command{ 13 | Use: "serve-dummy", 14 | Short: "Launches the Data Catalog server without any connections", 15 | RunE: func(cmd *cobra.Command, args []string) error { 16 | ctx := context.Background() 17 | cfg := config.GetConfig() 18 | 19 | // serve a http healthcheck endpoint 20 | go func() { 21 | err := datacatalogservice.ServeHTTPHealthCheck(ctx, cfg) 22 | if err != nil { 23 | logger.Errorf(ctx, "Unable to serve http", cfg.GetGrpcHostAddress(), err) 24 | } 25 | }() 26 | 27 | return datacatalogservice.Serve(ctx, cfg) 28 | }, 29 | } 30 | 31 | func init() { 32 | RootCmd.AddCommand(serveDummyCmd) 33 | } 34 | -------------------------------------------------------------------------------- /cmd/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/cmd/entrypoints" 5 | "github.com/golang/glog" 6 | ) 7 | 8 | func main() { 9 | glog.V(2).Info("Beginning Data Catalog") 10 | err := entrypoints.Execute() 11 | if err != nil { 12 | panic(err) 13 | } 14 | } 15 | -------------------------------------------------------------------------------- /datacatalog.json: -------------------------------------------------------------------------------- 1 | { 2 | "version": "1.0.38", 3 | "architecture": { 4 | "32bit": { 5 | "url": "https://github.com/flyteorg/datacatalog/releases/download/v1.0.38/datacatalog_v1.0.38_windows_i386.zip", 6 | "bin": [ 7 | "datacatalog.exe" 8 | ], 9 | "hash": "3aae4c7b2907e9c2d564363951f98ba7e22de20de5e5d9e10652303ef5c0f858" 10 | }, 11 | "64bit": { 12 | "url": "https://github.com/flyteorg/datacatalog/releases/download/v1.0.38/datacatalog_v1.0.38_windows_x86_64.zip", 13 | "bin": [ 14 | "datacatalog.exe" 15 | ], 16 | "hash": "f355deb945e0cc19e5a5ebfc589ab8d7f48363514490ec19f6ec7bbc9ba3a293" 17 | } 18 | }, 19 | "homepage": "https://godoc.org/github.com/lyft/datacatalog", 20 | "license": "Apache-2.0", 21 | "description": "datacatalog is the a memoization \u0026 lineage tracking service." 22 | } -------------------------------------------------------------------------------- /datacatalog_config.yaml: -------------------------------------------------------------------------------- 1 | # This is a sample configuration file. 2 | # Real configuration when running inside K8s (local or otherwise) lives in a ConfigMap 3 | # Look in the artifacts directory in the flyte repo for what's actually run 4 | logger: 5 | level: 5 6 | application: 7 | grpcPort: 8081 8 | httpPort: 8080 9 | grpcServerReflection: true 10 | datacatalog: 11 | storage-prefix: "metadata" 12 | metrics-scope: "datacatalog" 13 | profiler-port: 10254 14 | heartbeat-grace-period-multiplier: 3 15 | max-reservation-heartbeat: 10s 16 | storage: 17 | connection: 18 | access-key: minio 19 | auth-type: accesskey 20 | disable-ssl: true 21 | endpoint: http://localhost:9000 22 | region: my-region-here 23 | secret-key: miniostorage 24 | cache: 25 | max_size_mbs: 10 26 | target_gc_percent: 100 27 | container: my-container 28 | type: minio 29 | database: 30 | port: 5432 31 | username: postgres 32 | host: localhost 33 | dbname: datacatalog 34 | options: "sslmode=disable" 35 | log_level: 5 36 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/flyteorg/datacatalog 2 | 3 | go 1.19 4 | 5 | require ( 6 | github.com/Selvatico/go-mocket v1.0.7 7 | github.com/flyteorg/flyteidl v1.3.6 8 | github.com/flyteorg/flytestdlib v1.0.22 9 | github.com/gofrs/uuid v4.2.0+incompatible 10 | github.com/golang/glog v1.1.0 11 | github.com/golang/protobuf v1.5.3 12 | github.com/jackc/pgconn v1.10.1 13 | github.com/mitchellh/mapstructure v1.5.0 14 | github.com/spf13/cobra v1.4.0 15 | github.com/spf13/pflag v1.0.5 16 | github.com/stretchr/testify v1.8.4 17 | google.golang.org/grpc v1.56.1 18 | gorm.io/driver/postgres v1.2.3 19 | gorm.io/driver/sqlite v1.1.1 20 | gorm.io/gorm v1.22.4 21 | ) 22 | 23 | require ( 24 | cloud.google.com/go v0.110.0 // indirect 25 | cloud.google.com/go/compute v1.19.1 // indirect 26 | cloud.google.com/go/compute/metadata v0.2.3 // indirect 27 | cloud.google.com/go/iam v0.13.0 // indirect 28 | cloud.google.com/go/storage v1.28.1 // indirect 29 | github.com/Azure/azure-sdk-for-go v63.4.0+incompatible // indirect 30 | github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.1 // indirect 31 | github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2 // indirect 32 | github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v0.4.0 // indirect 33 | github.com/Azure/go-autorest v14.2.0+incompatible // indirect 34 | github.com/Azure/go-autorest/autorest v0.11.27 // indirect 35 | github.com/Azure/go-autorest/autorest/adal v0.9.18 // indirect 36 | github.com/Azure/go-autorest/autorest/date v0.3.0 // indirect 37 | github.com/Azure/go-autorest/logger v0.2.1 // indirect 38 | github.com/Azure/go-autorest/tracing v0.6.0 // indirect 39 | github.com/aws/aws-sdk-go v1.44.2 // indirect 40 | github.com/beorn7/perks v1.0.1 // indirect 41 | github.com/cespare/xxhash v1.1.0 // indirect 42 | github.com/cespare/xxhash/v2 v2.2.0 // indirect 43 | github.com/coocood/freecache v1.1.1 // indirect 44 | github.com/davecgh/go-spew v1.1.1 // indirect 45 | github.com/fatih/color v1.13.0 // indirect 46 | github.com/flyteorg/stow v0.3.7 // indirect 47 | github.com/fsnotify/fsnotify v1.5.1 // indirect 48 | github.com/ghodss/yaml v1.0.0 // indirect 49 | github.com/go-logr/logr v0.4.0 // indirect 50 | github.com/golang-jwt/jwt/v4 v4.4.1 // indirect 51 | github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect 52 | github.com/google/go-cmp v0.5.9 // indirect 53 | github.com/google/uuid v1.3.0 // indirect 54 | github.com/googleapis/enterprise-certificate-proxy v0.2.3 // indirect 55 | github.com/googleapis/gax-go/v2 v2.7.1 // indirect 56 | github.com/hashicorp/hcl v1.0.0 // indirect 57 | github.com/inconshreveable/mousetrap v1.0.0 // indirect 58 | github.com/jackc/chunkreader/v2 v2.0.1 // indirect 59 | github.com/jackc/pgio v1.0.0 // indirect 60 | github.com/jackc/pgpassfile v1.0.0 // indirect 61 | github.com/jackc/pgproto3/v2 v2.2.0 // indirect 62 | github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect 63 | github.com/jackc/pgtype v1.9.0 // indirect 64 | github.com/jackc/pgx/v4 v4.14.0 // indirect 65 | github.com/jinzhu/inflection v1.0.0 // indirect 66 | github.com/jinzhu/now v1.1.4 // indirect 67 | github.com/jmespath/go-jmespath v0.4.0 // indirect 68 | github.com/magiconair/properties v1.8.6 // indirect 69 | github.com/mattn/go-colorable v0.1.12 // indirect 70 | github.com/mattn/go-isatty v0.0.14 // indirect 71 | github.com/mattn/go-sqlite3 v1.14.0 // indirect 72 | github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect 73 | github.com/ncw/swift v1.0.53 // indirect 74 | github.com/pelletier/go-toml v1.9.4 // indirect 75 | github.com/pelletier/go-toml/v2 v2.0.0-beta.8 // indirect 76 | github.com/pkg/errors v0.9.1 // indirect 77 | github.com/pmezard/go-difflib v1.0.0 // indirect 78 | github.com/prometheus/client_golang v1.12.1 // indirect 79 | github.com/prometheus/client_model v0.2.0 // indirect 80 | github.com/prometheus/common v0.32.1 // indirect 81 | github.com/prometheus/procfs v0.7.3 // indirect 82 | github.com/sirupsen/logrus v1.7.0 // indirect 83 | github.com/spf13/afero v1.9.2 // indirect 84 | github.com/spf13/cast v1.4.1 // indirect 85 | github.com/spf13/jwalterweatherman v1.1.0 // indirect 86 | github.com/spf13/viper v1.11.0 // indirect 87 | github.com/stretchr/objx v0.5.0 // indirect 88 | github.com/subosito/gotenv v1.2.0 // indirect 89 | go.opencensus.io v0.24.0 // indirect 90 | golang.org/x/crypto v0.0.0-20220427172511-eb4f295cb31f // indirect 91 | golang.org/x/net v0.9.0 // indirect 92 | golang.org/x/oauth2 v0.7.0 // indirect 93 | golang.org/x/sys v0.7.0 // indirect 94 | golang.org/x/text v0.9.0 // indirect 95 | golang.org/x/time v0.1.0 // indirect 96 | golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect 97 | google.golang.org/api v0.114.0 // indirect 98 | google.golang.org/appengine v1.6.7 // indirect 99 | google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1 // indirect 100 | google.golang.org/protobuf v1.30.0 // indirect 101 | gopkg.in/ini.v1 v1.66.4 // indirect 102 | gopkg.in/yaml.v2 v2.4.0 // indirect 103 | gopkg.in/yaml.v3 v3.0.1 // indirect 104 | k8s.io/apimachinery v0.20.2 // indirect 105 | k8s.io/client-go v0.0.0-20210217172142-7279fc64d847 // indirect 106 | k8s.io/klog/v2 v2.5.0 // indirect 107 | ) 108 | -------------------------------------------------------------------------------- /pkg/common/filters.go: -------------------------------------------------------------------------------- 1 | package common 2 | 3 | // Common constants and types for Filtering 4 | const ( 5 | DefaultPageOffset = 0 6 | MaxPageLimit = 50 7 | ) 8 | 9 | // Common Entity types that can be used on any filters 10 | type Entity string 11 | 12 | const ( 13 | Artifact Entity = "Artifact" 14 | Dataset Entity = "Dataset" 15 | Partition Entity = "Partition" 16 | Tag Entity = "Tag" 17 | ) 18 | 19 | // Supported operators that can be used on filters 20 | type ComparisonOperator int 21 | 22 | const ( 23 | Equal ComparisonOperator = iota 24 | // Add more operators as needed, ie., gte, lte 25 | ) 26 | -------------------------------------------------------------------------------- /pkg/config/config.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/flyteorg/flytestdlib/config" 7 | ) 8 | 9 | const SectionKey = "application" 10 | 11 | //go:generate pflags Config 12 | 13 | type Config struct { 14 | GrpcPort int `json:"grpcPort" pflag:",On which grpc port to serve Catalog"` 15 | GrpcServerReflection bool `json:"grpcServerReflection" pflag:",Enable GRPC Server Reflection"` 16 | HTTPPort int `json:"httpPort" pflag:",On which http port to serve Catalog"` 17 | Secure bool `json:"secure" pflag:",Whether to run Catalog in secure mode or not"` 18 | ReadHeaderTimeoutSeconds int `json:"readHeaderTimeoutSeconds" pflag:",The amount of time allowed to read request headers."` 19 | } 20 | 21 | var defaultConfig = &Config{ 22 | GrpcPort: 8081, 23 | HTTPPort: 8080, 24 | GrpcServerReflection: true, 25 | // Set the HTTP timeout to avoid security vulnerabilities with expired, inactive connections: 26 | // https://deepsource.io/directory/analyzers/go/issues/GO-S2114 27 | // just shy of requestTimeoutUpperBound 28 | ReadHeaderTimeoutSeconds: 32, 29 | } 30 | var applicationConfig = config.MustRegisterSection(SectionKey, defaultConfig) 31 | 32 | func GetConfig() *Config { 33 | return applicationConfig.GetConfig().(*Config) 34 | } 35 | 36 | func SetConfig(c *Config) { 37 | if err := applicationConfig.SetConfig(c); err != nil { 38 | panic(err) 39 | } 40 | } 41 | 42 | func (c Config) GetGrpcHostAddress() string { 43 | return fmt.Sprintf(":%d", c.GrpcPort) 44 | } 45 | 46 | func (c Config) GetHTTPHostAddress() string { 47 | return fmt.Sprintf(":%d", c.HTTPPort) 48 | } 49 | 50 | func init() { 51 | SetConfig(&Config{}) 52 | } 53 | -------------------------------------------------------------------------------- /pkg/config/config_flags.go: -------------------------------------------------------------------------------- 1 | // Code generated by go generate; DO NOT EDIT. 2 | // This file was generated by robots. 3 | 4 | package config 5 | 6 | import ( 7 | "encoding/json" 8 | "reflect" 9 | 10 | "fmt" 11 | 12 | "github.com/spf13/pflag" 13 | ) 14 | 15 | // If v is a pointer, it will get its element value or the zero value of the element type. 16 | // If v is not a pointer, it will return it as is. 17 | func (Config) elemValueOrNil(v interface{}) interface{} { 18 | if t := reflect.TypeOf(v); t.Kind() == reflect.Ptr { 19 | if reflect.ValueOf(v).IsNil() { 20 | return reflect.Zero(t.Elem()).Interface() 21 | } else { 22 | return reflect.ValueOf(v).Interface() 23 | } 24 | } else if v == nil { 25 | return reflect.Zero(t).Interface() 26 | } 27 | 28 | return v 29 | } 30 | 31 | func (Config) mustJsonMarshal(v interface{}) string { 32 | raw, err := json.Marshal(v) 33 | if err != nil { 34 | panic(err) 35 | } 36 | 37 | return string(raw) 38 | } 39 | 40 | func (Config) mustMarshalJSON(v json.Marshaler) string { 41 | raw, err := v.MarshalJSON() 42 | if err != nil { 43 | panic(err) 44 | } 45 | 46 | return string(raw) 47 | } 48 | 49 | // GetPFlagSet will return strongly types pflags for all fields in Config and its nested types. The format of the 50 | // flags is json-name.json-sub-name... etc. 51 | func (cfg Config) GetPFlagSet(prefix string) *pflag.FlagSet { 52 | cmdFlags := pflag.NewFlagSet("Config", pflag.ExitOnError) 53 | cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "grpcPort"), defaultConfig.GrpcPort, "On which grpc port to serve Catalog") 54 | cmdFlags.Bool(fmt.Sprintf("%v%v", prefix, "grpcServerReflection"), defaultConfig.GrpcServerReflection, "Enable GRPC Server Reflection") 55 | cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "httpPort"), defaultConfig.HTTPPort, "On which http port to serve Catalog") 56 | cmdFlags.Bool(fmt.Sprintf("%v%v", prefix, "secure"), defaultConfig.Secure, "Whether to run Catalog in secure mode or not") 57 | cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "readHeaderTimeoutSeconds"), defaultConfig.ReadHeaderTimeoutSeconds, "The amount of time allowed to read request headers.") 58 | return cmdFlags 59 | } 60 | -------------------------------------------------------------------------------- /pkg/config/config_flags_test.go: -------------------------------------------------------------------------------- 1 | // Code generated by go generate; DO NOT EDIT. 2 | // This file was generated by robots. 3 | 4 | package config 5 | 6 | import ( 7 | "encoding/json" 8 | "fmt" 9 | "reflect" 10 | "strings" 11 | "testing" 12 | 13 | "github.com/mitchellh/mapstructure" 14 | "github.com/stretchr/testify/assert" 15 | ) 16 | 17 | var dereferencableKindsConfig = map[reflect.Kind]struct{}{ 18 | reflect.Array: {}, reflect.Chan: {}, reflect.Map: {}, reflect.Ptr: {}, reflect.Slice: {}, 19 | } 20 | 21 | // Checks if t is a kind that can be dereferenced to get its underlying type. 22 | func canGetElementConfig(t reflect.Kind) bool { 23 | _, exists := dereferencableKindsConfig[t] 24 | return exists 25 | } 26 | 27 | // This decoder hook tests types for json unmarshaling capability. If implemented, it uses json unmarshal to build the 28 | // object. Otherwise, it'll just pass on the original data. 29 | func jsonUnmarshalerHookConfig(_, to reflect.Type, data interface{}) (interface{}, error) { 30 | unmarshalerType := reflect.TypeOf((*json.Unmarshaler)(nil)).Elem() 31 | if to.Implements(unmarshalerType) || reflect.PtrTo(to).Implements(unmarshalerType) || 32 | (canGetElementConfig(to.Kind()) && to.Elem().Implements(unmarshalerType)) { 33 | 34 | raw, err := json.Marshal(data) 35 | if err != nil { 36 | fmt.Printf("Failed to marshal Data: %v. Error: %v. Skipping jsonUnmarshalHook", data, err) 37 | return data, nil 38 | } 39 | 40 | res := reflect.New(to).Interface() 41 | err = json.Unmarshal(raw, &res) 42 | if err != nil { 43 | fmt.Printf("Failed to umarshal Data: %v. Error: %v. Skipping jsonUnmarshalHook", data, err) 44 | return data, nil 45 | } 46 | 47 | return res, nil 48 | } 49 | 50 | return data, nil 51 | } 52 | 53 | func decode_Config(input, result interface{}) error { 54 | config := &mapstructure.DecoderConfig{ 55 | TagName: "json", 56 | WeaklyTypedInput: true, 57 | Result: result, 58 | DecodeHook: mapstructure.ComposeDecodeHookFunc( 59 | mapstructure.StringToTimeDurationHookFunc(), 60 | mapstructure.StringToSliceHookFunc(","), 61 | jsonUnmarshalerHookConfig, 62 | ), 63 | } 64 | 65 | decoder, err := mapstructure.NewDecoder(config) 66 | if err != nil { 67 | return err 68 | } 69 | 70 | return decoder.Decode(input) 71 | } 72 | 73 | func join_Config(arr interface{}, sep string) string { 74 | listValue := reflect.ValueOf(arr) 75 | strs := make([]string, 0, listValue.Len()) 76 | for i := 0; i < listValue.Len(); i++ { 77 | strs = append(strs, fmt.Sprintf("%v", listValue.Index(i))) 78 | } 79 | 80 | return strings.Join(strs, sep) 81 | } 82 | 83 | func testDecodeJson_Config(t *testing.T, val, result interface{}) { 84 | assert.NoError(t, decode_Config(val, result)) 85 | } 86 | 87 | func testDecodeRaw_Config(t *testing.T, vStringSlice, result interface{}) { 88 | assert.NoError(t, decode_Config(vStringSlice, result)) 89 | } 90 | 91 | func TestConfig_GetPFlagSet(t *testing.T) { 92 | val := Config{} 93 | cmdFlags := val.GetPFlagSet("") 94 | assert.True(t, cmdFlags.HasFlags()) 95 | } 96 | 97 | func TestConfig_SetFlags(t *testing.T) { 98 | actual := Config{} 99 | cmdFlags := actual.GetPFlagSet("") 100 | assert.True(t, cmdFlags.HasFlags()) 101 | 102 | t.Run("Test_grpcPort", func(t *testing.T) { 103 | 104 | t.Run("Override", func(t *testing.T) { 105 | testValue := "1" 106 | 107 | cmdFlags.Set("grpcPort", testValue) 108 | if vInt, err := cmdFlags.GetInt("grpcPort"); err == nil { 109 | testDecodeJson_Config(t, fmt.Sprintf("%v", vInt), &actual.GrpcPort) 110 | 111 | } else { 112 | assert.FailNow(t, err.Error()) 113 | } 114 | }) 115 | }) 116 | t.Run("Test_grpcServerReflection", func(t *testing.T) { 117 | 118 | t.Run("Override", func(t *testing.T) { 119 | testValue := "1" 120 | 121 | cmdFlags.Set("grpcServerReflection", testValue) 122 | if vBool, err := cmdFlags.GetBool("grpcServerReflection"); err == nil { 123 | testDecodeJson_Config(t, fmt.Sprintf("%v", vBool), &actual.GrpcServerReflection) 124 | 125 | } else { 126 | assert.FailNow(t, err.Error()) 127 | } 128 | }) 129 | }) 130 | t.Run("Test_httpPort", func(t *testing.T) { 131 | 132 | t.Run("Override", func(t *testing.T) { 133 | testValue := "1" 134 | 135 | cmdFlags.Set("httpPort", testValue) 136 | if vInt, err := cmdFlags.GetInt("httpPort"); err == nil { 137 | testDecodeJson_Config(t, fmt.Sprintf("%v", vInt), &actual.HTTPPort) 138 | 139 | } else { 140 | assert.FailNow(t, err.Error()) 141 | } 142 | }) 143 | }) 144 | t.Run("Test_secure", func(t *testing.T) { 145 | 146 | t.Run("Override", func(t *testing.T) { 147 | testValue := "1" 148 | 149 | cmdFlags.Set("secure", testValue) 150 | if vBool, err := cmdFlags.GetBool("secure"); err == nil { 151 | testDecodeJson_Config(t, fmt.Sprintf("%v", vBool), &actual.Secure) 152 | 153 | } else { 154 | assert.FailNow(t, err.Error()) 155 | } 156 | }) 157 | }) 158 | t.Run("Test_readHeaderTimeoutSeconds", func(t *testing.T) { 159 | 160 | t.Run("Override", func(t *testing.T) { 161 | testValue := "1" 162 | 163 | cmdFlags.Set("readHeaderTimeoutSeconds", testValue) 164 | if vInt, err := cmdFlags.GetInt("readHeaderTimeoutSeconds"); err == nil { 165 | testDecodeJson_Config(t, fmt.Sprintf("%v", vInt), &actual.ReadHeaderTimeoutSeconds) 166 | 167 | } else { 168 | assert.FailNow(t, err.Error()) 169 | } 170 | }) 171 | }) 172 | } 173 | -------------------------------------------------------------------------------- /pkg/errors/errors.go: -------------------------------------------------------------------------------- 1 | package errors 2 | 3 | import ( 4 | "fmt" 5 | "strings" 6 | 7 | "google.golang.org/grpc/codes" 8 | "google.golang.org/grpc/status" 9 | ) 10 | 11 | type DataCatalogError interface { 12 | Error() string 13 | Code() codes.Code 14 | GRPCStatus() *status.Status 15 | String() string 16 | } 17 | 18 | type dataCatalogErrorImpl struct { 19 | status *status.Status 20 | } 21 | 22 | func (e *dataCatalogErrorImpl) Error() string { 23 | return e.status.Message() 24 | } 25 | 26 | func (e *dataCatalogErrorImpl) Code() codes.Code { 27 | return e.status.Code() 28 | } 29 | 30 | func (e *dataCatalogErrorImpl) GRPCStatus() *status.Status { 31 | return e.status 32 | } 33 | 34 | func (e *dataCatalogErrorImpl) String() string { 35 | return fmt.Sprintf("status: %v", e.status) 36 | } 37 | 38 | func NewDataCatalogError(code codes.Code, message string) error { 39 | return &dataCatalogErrorImpl{ 40 | status: status.New(code, message), 41 | } 42 | } 43 | 44 | func NewDataCatalogErrorf(code codes.Code, format string, a ...interface{}) error { 45 | return NewDataCatalogError(code, fmt.Sprintf(format, a...)) 46 | } 47 | 48 | func NewCollectedErrors(code codes.Code, errors []error) error { 49 | errorCollection := make([]string, len(errors)) 50 | for idx, err := range errors { 51 | errorCollection[idx] = err.Error() 52 | } 53 | 54 | return NewDataCatalogError(code, strings.Join((errorCollection), ", ")) 55 | } 56 | 57 | func IsAlreadyExistsError(err error) bool { 58 | dcErr, ok := err.(DataCatalogError) 59 | return ok && dcErr.GRPCStatus().Code() == codes.AlreadyExists 60 | } 61 | 62 | func IsDoesNotExistError(err error) bool { 63 | dcErr, ok := err.(DataCatalogError) 64 | return ok && dcErr.GRPCStatus().Code() == codes.NotFound 65 | } 66 | -------------------------------------------------------------------------------- /pkg/errors/errors_test.go: -------------------------------------------------------------------------------- 1 | package errors 2 | 3 | import ( 4 | "testing" 5 | 6 | "fmt" 7 | 8 | "github.com/stretchr/testify/assert" 9 | "google.golang.org/grpc/codes" 10 | "google.golang.org/grpc/status" 11 | ) 12 | 13 | func TestErrorHelpers(t *testing.T) { 14 | alreadyExistsErr := NewDataCatalogError(codes.AlreadyExists, "already exists") 15 | notFoundErr := NewDataCatalogError(codes.NotFound, "not found") 16 | 17 | t.Run("TestAlreadyExists", func(t *testing.T) { 18 | assert.True(t, IsAlreadyExistsError(alreadyExistsErr)) 19 | assert.False(t, IsAlreadyExistsError(notFoundErr)) 20 | }) 21 | 22 | t.Run("TestNotFoundErr", func(t *testing.T) { 23 | assert.False(t, IsDoesNotExistError(alreadyExistsErr)) 24 | assert.True(t, IsDoesNotExistError(notFoundErr)) 25 | }) 26 | 27 | t.Run("TestCollectErrs", func(t *testing.T) { 28 | collectedErr := NewCollectedErrors(codes.InvalidArgument, []error{alreadyExistsErr, notFoundErr}) 29 | assert.EqualValues(t, status.Code(collectedErr), codes.InvalidArgument) 30 | assert.Equal(t, collectedErr.Error(), fmt.Sprintf("%s, %s", alreadyExistsErr.Error(), notFoundErr.Error())) 31 | }) 32 | } 33 | -------------------------------------------------------------------------------- /pkg/manager/impl/artifact_data_store.go: -------------------------------------------------------------------------------- 1 | package impl 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/errors" 7 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 8 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core" 9 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 10 | "github.com/flyteorg/flytestdlib/storage" 11 | "google.golang.org/grpc/codes" 12 | ) 13 | 14 | const artifactDataFile = "data.pb" 15 | 16 | // ArtifactDataStore stores and retrieves ArtifactData values in a data.pb 17 | type ArtifactDataStore interface { 18 | PutData(ctx context.Context, artifact *datacatalog.Artifact, data *datacatalog.ArtifactData) (storage.DataReference, error) 19 | GetData(ctx context.Context, dataModel models.ArtifactData) (*core.Literal, error) 20 | DeleteData(ctx context.Context, dataModel models.ArtifactData) error 21 | } 22 | 23 | type artifactDataStore struct { 24 | store *storage.DataStore 25 | storagePrefix storage.DataReference 26 | } 27 | 28 | func (m *artifactDataStore) getDataLocation(ctx context.Context, artifact *datacatalog.Artifact, data *datacatalog.ArtifactData) (storage.DataReference, error) { 29 | dataset := artifact.Dataset 30 | return m.store.ConstructReference(ctx, m.storagePrefix, dataset.Project, dataset.Domain, dataset.Name, dataset.Version, artifact.Id, data.Name, artifactDataFile) 31 | } 32 | 33 | // Store marshalled data in data.pb under the storage prefix 34 | func (m *artifactDataStore) PutData(ctx context.Context, artifact *datacatalog.Artifact, data *datacatalog.ArtifactData) (storage.DataReference, error) { 35 | dataLocation, err := m.getDataLocation(ctx, artifact, data) 36 | if err != nil { 37 | return "", errors.NewDataCatalogErrorf(codes.Internal, "Unable to generate data location %s, err %v", dataLocation.String(), err) 38 | } 39 | err = m.store.WriteProtobuf(ctx, dataLocation, storage.Options{}, data.Value) 40 | if err != nil { 41 | return "", errors.NewDataCatalogErrorf(codes.Internal, "Unable to store artifact data in location %s, err %v", dataLocation.String(), err) 42 | } 43 | 44 | return dataLocation, nil 45 | } 46 | 47 | // Retrieve the literal value of the ArtifactData from its specified location 48 | func (m *artifactDataStore) GetData(ctx context.Context, dataModel models.ArtifactData) (*core.Literal, error) { 49 | var value core.Literal 50 | err := m.store.ReadProtobuf(ctx, storage.DataReference(dataModel.Location), &value) 51 | if err != nil { 52 | return nil, errors.NewDataCatalogErrorf(codes.Internal, "Unable to read artifact data from location %s, err %v", dataModel.Location, err) 53 | } 54 | 55 | return &value, nil 56 | } 57 | 58 | // DeleteData removes the stored artifact data from the underlying blob storage 59 | func (m *artifactDataStore) DeleteData(ctx context.Context, dataModel models.ArtifactData) error { 60 | if err := m.store.Delete(ctx, storage.DataReference(dataModel.Location)); err != nil { 61 | return errors.NewDataCatalogErrorf(codes.Internal, "Unable to delete artifact data in location %s, err %v", dataModel.Location, err) 62 | } 63 | 64 | return nil 65 | } 66 | 67 | func NewArtifactDataStore(store *storage.DataStore, storagePrefix storage.DataReference) ArtifactDataStore { 68 | return &artifactDataStore{ 69 | store: store, 70 | storagePrefix: storagePrefix, 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /pkg/manager/impl/tag_manager.go: -------------------------------------------------------------------------------- 1 | package impl 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/manager/impl/validators" 8 | "github.com/flyteorg/datacatalog/pkg/manager/interfaces" 9 | "github.com/flyteorg/datacatalog/pkg/repositories" 10 | 11 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 12 | "github.com/flyteorg/datacatalog/pkg/repositories/transformers" 13 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 14 | 15 | "github.com/flyteorg/datacatalog/pkg/errors" 16 | "github.com/flyteorg/flytestdlib/contextutils" 17 | "github.com/flyteorg/flytestdlib/logger" 18 | "github.com/flyteorg/flytestdlib/promutils" 19 | "github.com/flyteorg/flytestdlib/promutils/labeled" 20 | "github.com/flyteorg/flytestdlib/storage" 21 | ) 22 | 23 | type tagMetrics struct { 24 | scope promutils.Scope 25 | createResponseTime labeled.StopWatch 26 | addTagSuccessCounter labeled.Counter 27 | addTagFailureCounter labeled.Counter 28 | validationErrorCounter labeled.Counter 29 | alreadyExistsCounter labeled.Counter 30 | } 31 | 32 | type tagManager struct { 33 | repo repositories.RepositoryInterface 34 | store *storage.DataStore 35 | systemMetrics tagMetrics 36 | } 37 | 38 | func (m *tagManager) AddTag(ctx context.Context, request *datacatalog.AddTagRequest) (*datacatalog.AddTagResponse, error) { 39 | timer := m.systemMetrics.createResponseTime.Start(ctx) 40 | defer timer.Stop() 41 | 42 | if err := validators.ValidateTag(request.Tag); err != nil { 43 | logger.Warnf(ctx, "Invalid get tag request %+v err: %v", request, err) 44 | m.systemMetrics.validationErrorCounter.Inc(ctx) 45 | return nil, err 46 | } 47 | 48 | // verify the artifact and dataset exists before adding a tag to it 49 | datasetID := request.Tag.Dataset 50 | ctx = contextutils.WithProjectDomain(ctx, datasetID.Project, datasetID.Domain) 51 | 52 | datasetKey := transformers.FromDatasetID(datasetID) 53 | dataset, err := m.repo.DatasetRepo().Get(ctx, datasetKey) 54 | if err != nil { 55 | m.systemMetrics.addTagFailureCounter.Inc(ctx) 56 | return nil, err 57 | } 58 | 59 | artifactKey := transformers.ToArtifactKey(datasetID, request.Tag.ArtifactId) 60 | _, err = m.repo.ArtifactRepo().Get(ctx, artifactKey) 61 | if err != nil { 62 | m.systemMetrics.addTagFailureCounter.Inc(ctx) 63 | return nil, err 64 | } 65 | 66 | tagKey := transformers.ToTagKey(datasetID, request.Tag.Name) 67 | err = m.repo.TagRepo().Create(ctx, models.Tag{ 68 | TagKey: tagKey, 69 | ArtifactID: request.Tag.ArtifactId, 70 | DatasetUUID: dataset.UUID, 71 | }) 72 | if err != nil { 73 | if errors.IsAlreadyExistsError(err) { 74 | logger.Warnf(ctx, "Tag already exists key: %+v, err %v", request, err) 75 | m.systemMetrics.alreadyExistsCounter.Inc(ctx) 76 | } else { 77 | logger.Errorf(ctx, "Failed to tag artifact: %+v err: %v", request, err) 78 | m.systemMetrics.addTagFailureCounter.Inc(ctx) 79 | } 80 | 81 | return nil, err 82 | } 83 | 84 | m.systemMetrics.addTagSuccessCounter.Inc(ctx) 85 | return &datacatalog.AddTagResponse{}, nil 86 | } 87 | 88 | func NewTagManager(repo repositories.RepositoryInterface, store *storage.DataStore, tagScope promutils.Scope) interfaces.TagManager { 89 | systemMetrics := tagMetrics{ 90 | scope: tagScope, 91 | createResponseTime: labeled.NewStopWatch("create_duration", "The duration of the add tag calls.", time.Millisecond, tagScope, labeled.EmitUnlabeledMetric), 92 | addTagSuccessCounter: labeled.NewCounter("create_success_count", "The number of times an artifact was tagged successfully", tagScope, labeled.EmitUnlabeledMetric), 93 | addTagFailureCounter: labeled.NewCounter("create_failure_count", "The number of times we failed to tag an artifact", tagScope, labeled.EmitUnlabeledMetric), 94 | validationErrorCounter: labeled.NewCounter("validation_failed_count", "The number of times we failed validate a tag", tagScope, labeled.EmitUnlabeledMetric), 95 | alreadyExistsCounter: labeled.NewCounter("already_exists_count", "The number of times an tag already exists", tagScope, labeled.EmitUnlabeledMetric), 96 | } 97 | 98 | return &tagManager{ 99 | repo: repo, 100 | store: store, 101 | systemMetrics: systemMetrics, 102 | } 103 | } 104 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/artifact_validator.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/common" 7 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 8 | ) 9 | 10 | const ( 11 | artifactID = "artifactID" 12 | artifactDataEntity = "artifactData" 13 | artifactEntity = "artifact" 14 | ) 15 | 16 | func ValidateGetArtifactRequest(request *datacatalog.GetArtifactRequest) error { 17 | if request.QueryHandle == nil { 18 | return NewMissingArgumentError(fmt.Sprintf("one of %s/%s", artifactID, tagName)) 19 | } 20 | 21 | switch request.QueryHandle.(type) { 22 | case *datacatalog.GetArtifactRequest_ArtifactId: 23 | if request.Dataset != nil { 24 | err := ValidateDatasetID(request.Dataset) 25 | if err != nil { 26 | return err 27 | } 28 | } 29 | 30 | if err := ValidateEmptyStringField(request.GetArtifactId(), artifactID); err != nil { 31 | return err 32 | } 33 | case *datacatalog.GetArtifactRequest_TagName: 34 | if err := ValidateDatasetID(request.Dataset); err != nil { 35 | return err 36 | } 37 | 38 | if err := ValidateEmptyStringField(request.GetTagName(), tagName); err != nil { 39 | return err 40 | } 41 | default: 42 | return NewInvalidArgumentError("QueryHandle", "invalid type") 43 | } 44 | 45 | return nil 46 | } 47 | 48 | func ValidateEmptyArtifactData(artifactData []*datacatalog.ArtifactData) error { 49 | if len(artifactData) == 0 { 50 | return NewMissingArgumentError(artifactDataEntity) 51 | } 52 | 53 | return nil 54 | } 55 | 56 | func ValidateArtifact(artifact *datacatalog.Artifact) error { 57 | if artifact == nil { 58 | return NewMissingArgumentError(artifactEntity) 59 | } 60 | 61 | if err := ValidateDatasetID(artifact.Dataset); err != nil { 62 | return err 63 | } 64 | 65 | if err := ValidateEmptyStringField(artifact.Id, artifactID); err != nil { 66 | return err 67 | } 68 | 69 | if err := ValidateEmptyArtifactData(artifact.Data); err != nil { 70 | return err 71 | } 72 | 73 | return nil 74 | } 75 | 76 | // Validate the list request and format the request with proper defaults if not provided 77 | func ValidateListArtifactRequest(request *datacatalog.ListArtifactsRequest) error { 78 | if err := ValidateDatasetID(request.Dataset); err != nil { 79 | return err 80 | } 81 | 82 | if err := ValidateArtifactFilterTypes(request.Filter.GetFilters()); err != nil { 83 | return err 84 | } 85 | 86 | if request.Pagination != nil { 87 | err := ValidatePagination(request.Pagination) 88 | if err != nil { 89 | return err 90 | } 91 | } 92 | 93 | return nil 94 | } 95 | 96 | // Artifacts cannot be filtered across Datasets 97 | func ValidateArtifactFilterTypes(filters []*datacatalog.SinglePropertyFilter) error { 98 | for _, filter := range filters { 99 | if filter.GetDatasetFilter() != nil { 100 | return NewInvalidFilterError(common.Artifact, common.Dataset) 101 | } 102 | } 103 | return nil 104 | } 105 | 106 | func ValidateUpdateArtifactRequest(request *datacatalog.UpdateArtifactRequest) error { 107 | if request.QueryHandle == nil { 108 | return NewMissingArgumentError(fmt.Sprintf("one of %s/%s", artifactID, tagName)) 109 | } 110 | 111 | switch request.QueryHandle.(type) { 112 | case *datacatalog.UpdateArtifactRequest_ArtifactId: 113 | if request.Dataset != nil { 114 | err := ValidateDatasetID(request.Dataset) 115 | if err != nil { 116 | return err 117 | } 118 | } 119 | 120 | if err := ValidateEmptyStringField(request.GetArtifactId(), artifactID); err != nil { 121 | return err 122 | } 123 | case *datacatalog.UpdateArtifactRequest_TagName: 124 | if err := ValidateDatasetID(request.Dataset); err != nil { 125 | return err 126 | } 127 | 128 | if err := ValidateEmptyStringField(request.GetTagName(), tagName); err != nil { 129 | return err 130 | } 131 | default: 132 | return NewInvalidArgumentError("QueryHandle", "invalid type") 133 | } 134 | 135 | if err := ValidateEmptyArtifactData(request.Data); err != nil { 136 | return err 137 | } 138 | 139 | return nil 140 | } 141 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/common.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | func ValidateEmptyStringField(field, fieldName string) error { 4 | if field == "" { 5 | return NewMissingArgumentError(fieldName) 6 | } 7 | return nil 8 | } 9 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/dataset_validator.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/common" 5 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 6 | ) 7 | 8 | const ( 9 | datasetEntity = "dataset" 10 | datasetProject = "project" 11 | datasetDomain = "domain" 12 | datasetName = "name" 13 | datasetVersion = "version" 14 | ) 15 | 16 | // Validate that the DatasetID has all the fields filled 17 | func ValidateDatasetID(ds *datacatalog.DatasetID) error { 18 | if ds == nil { 19 | return NewMissingArgumentError(datasetEntity) 20 | } 21 | if err := ValidateEmptyStringField(ds.Project, datasetProject); err != nil { 22 | return err 23 | } 24 | if err := ValidateEmptyStringField(ds.Domain, datasetDomain); err != nil { 25 | return err 26 | } 27 | if err := ValidateEmptyStringField(ds.Name, datasetName); err != nil { 28 | return err 29 | } 30 | if err := ValidateEmptyStringField(ds.Version, datasetVersion); err != nil { 31 | return err 32 | } 33 | return nil 34 | } 35 | 36 | // Ensure list Datasets request is properly constructed 37 | func ValidateListDatasetsRequest(request *datacatalog.ListDatasetsRequest) error { 38 | if request.Pagination != nil { 39 | err := ValidatePagination(request.Pagination) 40 | if err != nil { 41 | return err 42 | } 43 | } 44 | 45 | // Datasets cannot be filtered by tag, partitions or artifacts 46 | for _, filter := range request.Filter.GetFilters() { 47 | if filter.GetTagFilter() != nil { 48 | return NewInvalidFilterError(common.Dataset, common.Tag) 49 | } else if filter.GetPartitionFilter() != nil { 50 | return NewInvalidFilterError(common.Dataset, common.Partition) 51 | } else if filter.GetArtifactFilter() != nil { 52 | return NewInvalidFilterError(common.Dataset, common.Artifact) 53 | } 54 | } 55 | return nil 56 | } 57 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/errors.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/errors" 7 | 8 | "github.com/flyteorg/datacatalog/pkg/common" 9 | "google.golang.org/grpc/codes" 10 | ) 11 | 12 | const missingFieldFormat = "missing %s" 13 | const invalidArgFormat = "invalid value for %s, value:[%s]" 14 | const invalidFilterFormat = "%s cannot be filtered by %s properties" 15 | 16 | func NewMissingArgumentError(field string) error { 17 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, fmt.Sprintf(missingFieldFormat, field)) 18 | } 19 | 20 | func NewInvalidArgumentError(field string, value string) error { 21 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, fmt.Sprintf(invalidArgFormat, field, value)) 22 | } 23 | 24 | func NewInvalidFilterError(entity common.Entity, propertyEntity common.Entity) error { 25 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, fmt.Sprintf(invalidFilterFormat, entity, propertyEntity)) 26 | } 27 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/pagination_validator.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | import ( 4 | "strconv" 5 | "strings" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/errors" 8 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | "google.golang.org/grpc/codes" 10 | ) 11 | 12 | // The token is a string that should be opaque to the client 13 | // It represents the offset as an integer encoded as a string, 14 | // but in the future it can be a string that encodes anything 15 | func ValidateToken(token string) error { 16 | // if the token is empty, that is still valid input since it is optional 17 | if len(strings.Trim(token, " ")) == 0 { 18 | return nil 19 | } 20 | _, err := strconv.ParseUint(token, 10, 32) 21 | if err != nil { 22 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, "Invalid token value: %s", token) 23 | } 24 | return nil 25 | } 26 | 27 | // Validate the pagination options and set default limits 28 | func ValidatePagination(options *datacatalog.PaginationOptions) error { 29 | err := ValidateToken(options.Token) 30 | if err != nil { 31 | return err 32 | } 33 | 34 | if options.SortKey != datacatalog.PaginationOptions_CREATION_TIME { 35 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, "Invalid sort key %v", options.SortKey) 36 | } 37 | 38 | if options.SortOrder != datacatalog.PaginationOptions_ASCENDING && 39 | options.SortOrder != datacatalog.PaginationOptions_DESCENDING { 40 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, "Invalid sort order %v", options.SortOrder) 41 | } 42 | 43 | return nil 44 | } 45 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/partition_validator.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/errors" 7 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 8 | "google.golang.org/grpc/codes" 9 | ) 10 | 11 | const ( 12 | partitionKeyName = "partitionKey" 13 | partitionValueName = "partitionValue" 14 | ) 15 | 16 | func ValidatePartitions(datasetPartitionKeys []string, artifactPartitions []*datacatalog.Partition) error { 17 | if len(datasetPartitionKeys) != len(artifactPartitions) { 18 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, "Partition key mismatch, dataset keys: %+v, artifact Partitions: %+v", datasetPartitionKeys, artifactPartitions) 19 | } 20 | 21 | // Not all datasets need to be partitioned 22 | if len(datasetPartitionKeys) == 0 { 23 | return nil 24 | } 25 | 26 | // compare the contents of the datasetkeys and artifact keys 27 | partitionErrors := make([]error, 0) 28 | keyMismatch := false 29 | 30 | partitionKeyMatches := make(map[string]bool, len(artifactPartitions)) 31 | for _, datasetPartitionKey := range datasetPartitionKeys { 32 | partitionKeyMatches[datasetPartitionKey] = false 33 | } 34 | 35 | for idx, artifactPartition := range artifactPartitions { 36 | if artifactPartition == nil { 37 | partitionErrors = append(partitionErrors, NewMissingArgumentError(fmt.Sprintf("%v[%v]", partitionKeyName, idx))) 38 | continue 39 | } 40 | 41 | if err := ValidateEmptyStringField(partitionKeyName, artifactPartition.Key); err != nil { 42 | partitionErrors = append(partitionErrors, NewMissingArgumentError(fmt.Sprintf("%v[%v]", partitionKeyName, idx))) 43 | } else if err := ValidateEmptyStringField(partitionValueName, artifactPartition.Value); err != nil { 44 | partitionErrors = append(partitionErrors, NewMissingArgumentError(fmt.Sprintf("%v[%v]", partitionValueName, idx))) 45 | } else { 46 | _, ok := partitionKeyMatches[artifactPartition.Key] 47 | 48 | if ok { 49 | partitionKeyMatches[artifactPartition.Key] = true 50 | } else { 51 | keyMismatch = true 52 | } 53 | } 54 | } 55 | 56 | if keyMismatch { 57 | partitionErrors = append(partitionErrors, errors.NewDataCatalogErrorf(codes.InvalidArgument, "Artifact partition assignment does not match dataset partition keys: %v", partitionKeyMatches)) 58 | } 59 | 60 | if len(partitionErrors) > 0 { 61 | return errors.NewCollectedErrors(codes.InvalidArgument, partitionErrors) 62 | } 63 | 64 | return nil 65 | } 66 | 67 | // Validate that the partition keys are unique strings 68 | func ValidateUniquePartitionKeys(partitionKeys []string) error { 69 | invalidPartitionKeys := false 70 | partitionKeySet := make(map[string]uint8, len(partitionKeys)) 71 | for _, partitionKey := range partitionKeys { 72 | partitionKeySet[partitionKey]++ 73 | if partitionKeySet[partitionKey] > 1 { 74 | invalidPartitionKeys = true 75 | } 76 | } 77 | 78 | if invalidPartitionKeys { 79 | return NewInvalidArgumentError(partitionKeyName, fmt.Sprintf("Keys are not unique, occurrence count: %+v", partitionKeySet)) 80 | } 81 | 82 | return nil 83 | } 84 | -------------------------------------------------------------------------------- /pkg/manager/impl/validators/tag_validator.go: -------------------------------------------------------------------------------- 1 | package validators 2 | 3 | import ( 4 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 5 | ) 6 | 7 | const ( 8 | tagName = "tagName" 9 | tagEntity = "tag" 10 | ) 11 | 12 | func ValidateTag(tag *datacatalog.Tag) error { 13 | if tag == nil { 14 | return NewMissingArgumentError(tagEntity) 15 | } 16 | if err := ValidateDatasetID(tag.Dataset); err != nil { 17 | return err 18 | } 19 | 20 | if err := ValidateEmptyStringField(tag.Name, tagName); err != nil { 21 | return err 22 | } 23 | 24 | if err := ValidateEmptyStringField(tag.ArtifactId, artifactID); err != nil { 25 | return err 26 | } 27 | return nil 28 | } 29 | -------------------------------------------------------------------------------- /pkg/manager/interfaces/artifact.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | idl_datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | ) 8 | 9 | //go:generate mockery -all -output=../mocks -case=underscore 10 | 11 | type ArtifactManager interface { 12 | CreateArtifact(ctx context.Context, request *idl_datacatalog.CreateArtifactRequest) (*idl_datacatalog.CreateArtifactResponse, error) 13 | GetArtifact(ctx context.Context, request *idl_datacatalog.GetArtifactRequest) (*idl_datacatalog.GetArtifactResponse, error) 14 | ListArtifacts(ctx context.Context, request *idl_datacatalog.ListArtifactsRequest) (*idl_datacatalog.ListArtifactsResponse, error) 15 | UpdateArtifact(ctx context.Context, request *idl_datacatalog.UpdateArtifactRequest) (*idl_datacatalog.UpdateArtifactResponse, error) 16 | } 17 | -------------------------------------------------------------------------------- /pkg/manager/interfaces/dataset.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | idl_datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | ) 8 | 9 | type DatasetManager interface { 10 | CreateDataset(ctx context.Context, request *idl_datacatalog.CreateDatasetRequest) (*idl_datacatalog.CreateDatasetResponse, error) 11 | GetDataset(ctx context.Context, request *idl_datacatalog.GetDatasetRequest) (*idl_datacatalog.GetDatasetResponse, error) 12 | ListDatasets(ctx context.Context, request *idl_datacatalog.ListDatasetsRequest) (*idl_datacatalog.ListDatasetsResponse, error) 13 | } 14 | -------------------------------------------------------------------------------- /pkg/manager/interfaces/reservation.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | ) 8 | 9 | // ReservationManager is the interface to handle reservation requests. 10 | // You can find more details about the APIs in datacatalog service proto 11 | // in flyteidl 12 | type ReservationManager interface { 13 | GetOrExtendReservation(context.Context, *datacatalog.GetOrExtendReservationRequest) (*datacatalog.GetOrExtendReservationResponse, error) 14 | ReleaseReservation(context.Context, *datacatalog.ReleaseReservationRequest) (*datacatalog.ReleaseReservationResponse, error) 15 | } 16 | -------------------------------------------------------------------------------- /pkg/manager/interfaces/tag.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | ) 8 | 9 | type TagManager interface { 10 | AddTag(ctx context.Context, request *datacatalog.AddTagRequest) (*datacatalog.AddTagResponse, error) 11 | } 12 | -------------------------------------------------------------------------------- /pkg/manager/mocks/dataset_manager.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | 10 | mock "github.com/stretchr/testify/mock" 11 | ) 12 | 13 | // DatasetManager is an autogenerated mock type for the DatasetManager type 14 | type DatasetManager struct { 15 | mock.Mock 16 | } 17 | 18 | type DatasetManager_CreateDataset struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m DatasetManager_CreateDataset) Return(_a0 *datacatalog.CreateDatasetResponse, _a1 error) *DatasetManager_CreateDataset { 23 | return &DatasetManager_CreateDataset{Call: _m.Call.Return(_a0, _a1)} 24 | } 25 | 26 | func (_m *DatasetManager) OnCreateDataset(ctx context.Context, request *datacatalog.CreateDatasetRequest) *DatasetManager_CreateDataset { 27 | c_call := _m.On("CreateDataset", ctx, request) 28 | return &DatasetManager_CreateDataset{Call: c_call} 29 | } 30 | 31 | func (_m *DatasetManager) OnCreateDatasetMatch(matchers ...interface{}) *DatasetManager_CreateDataset { 32 | c_call := _m.On("CreateDataset", matchers...) 33 | return &DatasetManager_CreateDataset{Call: c_call} 34 | } 35 | 36 | // CreateDataset provides a mock function with given fields: ctx, request 37 | func (_m *DatasetManager) CreateDataset(ctx context.Context, request *datacatalog.CreateDatasetRequest) (*datacatalog.CreateDatasetResponse, error) { 38 | ret := _m.Called(ctx, request) 39 | 40 | var r0 *datacatalog.CreateDatasetResponse 41 | if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.CreateDatasetRequest) *datacatalog.CreateDatasetResponse); ok { 42 | r0 = rf(ctx, request) 43 | } else { 44 | if ret.Get(0) != nil { 45 | r0 = ret.Get(0).(*datacatalog.CreateDatasetResponse) 46 | } 47 | } 48 | 49 | var r1 error 50 | if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.CreateDatasetRequest) error); ok { 51 | r1 = rf(ctx, request) 52 | } else { 53 | r1 = ret.Error(1) 54 | } 55 | 56 | return r0, r1 57 | } 58 | 59 | type DatasetManager_GetDataset struct { 60 | *mock.Call 61 | } 62 | 63 | func (_m DatasetManager_GetDataset) Return(_a0 *datacatalog.GetDatasetResponse, _a1 error) *DatasetManager_GetDataset { 64 | return &DatasetManager_GetDataset{Call: _m.Call.Return(_a0, _a1)} 65 | } 66 | 67 | func (_m *DatasetManager) OnGetDataset(ctx context.Context, request *datacatalog.GetDatasetRequest) *DatasetManager_GetDataset { 68 | c_call := _m.On("GetDataset", ctx, request) 69 | return &DatasetManager_GetDataset{Call: c_call} 70 | } 71 | 72 | func (_m *DatasetManager) OnGetDatasetMatch(matchers ...interface{}) *DatasetManager_GetDataset { 73 | c_call := _m.On("GetDataset", matchers...) 74 | return &DatasetManager_GetDataset{Call: c_call} 75 | } 76 | 77 | // GetDataset provides a mock function with given fields: ctx, request 78 | func (_m *DatasetManager) GetDataset(ctx context.Context, request *datacatalog.GetDatasetRequest) (*datacatalog.GetDatasetResponse, error) { 79 | ret := _m.Called(ctx, request) 80 | 81 | var r0 *datacatalog.GetDatasetResponse 82 | if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.GetDatasetRequest) *datacatalog.GetDatasetResponse); ok { 83 | r0 = rf(ctx, request) 84 | } else { 85 | if ret.Get(0) != nil { 86 | r0 = ret.Get(0).(*datacatalog.GetDatasetResponse) 87 | } 88 | } 89 | 90 | var r1 error 91 | if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.GetDatasetRequest) error); ok { 92 | r1 = rf(ctx, request) 93 | } else { 94 | r1 = ret.Error(1) 95 | } 96 | 97 | return r0, r1 98 | } 99 | 100 | type DatasetManager_ListDatasets struct { 101 | *mock.Call 102 | } 103 | 104 | func (_m DatasetManager_ListDatasets) Return(_a0 *datacatalog.ListDatasetsResponse, _a1 error) *DatasetManager_ListDatasets { 105 | return &DatasetManager_ListDatasets{Call: _m.Call.Return(_a0, _a1)} 106 | } 107 | 108 | func (_m *DatasetManager) OnListDatasets(ctx context.Context, request *datacatalog.ListDatasetsRequest) *DatasetManager_ListDatasets { 109 | c_call := _m.On("ListDatasets", ctx, request) 110 | return &DatasetManager_ListDatasets{Call: c_call} 111 | } 112 | 113 | func (_m *DatasetManager) OnListDatasetsMatch(matchers ...interface{}) *DatasetManager_ListDatasets { 114 | c_call := _m.On("ListDatasets", matchers...) 115 | return &DatasetManager_ListDatasets{Call: c_call} 116 | } 117 | 118 | // ListDatasets provides a mock function with given fields: ctx, request 119 | func (_m *DatasetManager) ListDatasets(ctx context.Context, request *datacatalog.ListDatasetsRequest) (*datacatalog.ListDatasetsResponse, error) { 120 | ret := _m.Called(ctx, request) 121 | 122 | var r0 *datacatalog.ListDatasetsResponse 123 | if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.ListDatasetsRequest) *datacatalog.ListDatasetsResponse); ok { 124 | r0 = rf(ctx, request) 125 | } else { 126 | if ret.Get(0) != nil { 127 | r0 = ret.Get(0).(*datacatalog.ListDatasetsResponse) 128 | } 129 | } 130 | 131 | var r1 error 132 | if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.ListDatasetsRequest) error); ok { 133 | r1 = rf(ctx, request) 134 | } else { 135 | r1 = ret.Error(1) 136 | } 137 | 138 | return r0, r1 139 | } 140 | -------------------------------------------------------------------------------- /pkg/manager/mocks/reservation_manager.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | 10 | mock "github.com/stretchr/testify/mock" 11 | ) 12 | 13 | // ReservationManager is an autogenerated mock type for the ReservationManager type 14 | type ReservationManager struct { 15 | mock.Mock 16 | } 17 | 18 | type ReservationManager_GetOrExtendReservation struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m ReservationManager_GetOrExtendReservation) Return(_a0 *datacatalog.GetOrExtendReservationResponse, _a1 error) *ReservationManager_GetOrExtendReservation { 23 | return &ReservationManager_GetOrExtendReservation{Call: _m.Call.Return(_a0, _a1)} 24 | } 25 | 26 | func (_m *ReservationManager) OnGetOrExtendReservation(_a0 context.Context, _a1 *datacatalog.GetOrExtendReservationRequest) *ReservationManager_GetOrExtendReservation { 27 | c_call := _m.On("GetOrExtendReservation", _a0, _a1) 28 | return &ReservationManager_GetOrExtendReservation{Call: c_call} 29 | } 30 | 31 | func (_m *ReservationManager) OnGetOrExtendReservationMatch(matchers ...interface{}) *ReservationManager_GetOrExtendReservation { 32 | c_call := _m.On("GetOrExtendReservation", matchers...) 33 | return &ReservationManager_GetOrExtendReservation{Call: c_call} 34 | } 35 | 36 | // GetOrExtendReservation provides a mock function with given fields: _a0, _a1 37 | func (_m *ReservationManager) GetOrExtendReservation(_a0 context.Context, _a1 *datacatalog.GetOrExtendReservationRequest) (*datacatalog.GetOrExtendReservationResponse, error) { 38 | ret := _m.Called(_a0, _a1) 39 | 40 | var r0 *datacatalog.GetOrExtendReservationResponse 41 | if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.GetOrExtendReservationRequest) *datacatalog.GetOrExtendReservationResponse); ok { 42 | r0 = rf(_a0, _a1) 43 | } else { 44 | if ret.Get(0) != nil { 45 | r0 = ret.Get(0).(*datacatalog.GetOrExtendReservationResponse) 46 | } 47 | } 48 | 49 | var r1 error 50 | if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.GetOrExtendReservationRequest) error); ok { 51 | r1 = rf(_a0, _a1) 52 | } else { 53 | r1 = ret.Error(1) 54 | } 55 | 56 | return r0, r1 57 | } 58 | 59 | type ReservationManager_ReleaseReservation struct { 60 | *mock.Call 61 | } 62 | 63 | func (_m ReservationManager_ReleaseReservation) Return(_a0 *datacatalog.ReleaseReservationResponse, _a1 error) *ReservationManager_ReleaseReservation { 64 | return &ReservationManager_ReleaseReservation{Call: _m.Call.Return(_a0, _a1)} 65 | } 66 | 67 | func (_m *ReservationManager) OnReleaseReservation(_a0 context.Context, _a1 *datacatalog.ReleaseReservationRequest) *ReservationManager_ReleaseReservation { 68 | c_call := _m.On("ReleaseReservation", _a0, _a1) 69 | return &ReservationManager_ReleaseReservation{Call: c_call} 70 | } 71 | 72 | func (_m *ReservationManager) OnReleaseReservationMatch(matchers ...interface{}) *ReservationManager_ReleaseReservation { 73 | c_call := _m.On("ReleaseReservation", matchers...) 74 | return &ReservationManager_ReleaseReservation{Call: c_call} 75 | } 76 | 77 | // ReleaseReservation provides a mock function with given fields: _a0, _a1 78 | func (_m *ReservationManager) ReleaseReservation(_a0 context.Context, _a1 *datacatalog.ReleaseReservationRequest) (*datacatalog.ReleaseReservationResponse, error) { 79 | ret := _m.Called(_a0, _a1) 80 | 81 | var r0 *datacatalog.ReleaseReservationResponse 82 | if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.ReleaseReservationRequest) *datacatalog.ReleaseReservationResponse); ok { 83 | r0 = rf(_a0, _a1) 84 | } else { 85 | if ret.Get(0) != nil { 86 | r0 = ret.Get(0).(*datacatalog.ReleaseReservationResponse) 87 | } 88 | } 89 | 90 | var r1 error 91 | if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.ReleaseReservationRequest) error); ok { 92 | r1 = rf(_a0, _a1) 93 | } else { 94 | r1 = ret.Error(1) 95 | } 96 | 97 | return r0, r1 98 | } 99 | -------------------------------------------------------------------------------- /pkg/manager/mocks/tag_manager.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | 10 | mock "github.com/stretchr/testify/mock" 11 | ) 12 | 13 | // TagManager is an autogenerated mock type for the TagManager type 14 | type TagManager struct { 15 | mock.Mock 16 | } 17 | 18 | type TagManager_AddTag struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m TagManager_AddTag) Return(_a0 *datacatalog.AddTagResponse, _a1 error) *TagManager_AddTag { 23 | return &TagManager_AddTag{Call: _m.Call.Return(_a0, _a1)} 24 | } 25 | 26 | func (_m *TagManager) OnAddTag(ctx context.Context, request *datacatalog.AddTagRequest) *TagManager_AddTag { 27 | c_call := _m.On("AddTag", ctx, request) 28 | return &TagManager_AddTag{Call: c_call} 29 | } 30 | 31 | func (_m *TagManager) OnAddTagMatch(matchers ...interface{}) *TagManager_AddTag { 32 | c_call := _m.On("AddTag", matchers...) 33 | return &TagManager_AddTag{Call: c_call} 34 | } 35 | 36 | // AddTag provides a mock function with given fields: ctx, request 37 | func (_m *TagManager) AddTag(ctx context.Context, request *datacatalog.AddTagRequest) (*datacatalog.AddTagResponse, error) { 38 | ret := _m.Called(ctx, request) 39 | 40 | var r0 *datacatalog.AddTagResponse 41 | if rf, ok := ret.Get(0).(func(context.Context, *datacatalog.AddTagRequest) *datacatalog.AddTagResponse); ok { 42 | r0 = rf(ctx, request) 43 | } else { 44 | if ret.Get(0) != nil { 45 | r0 = ret.Get(0).(*datacatalog.AddTagResponse) 46 | } 47 | } 48 | 49 | var r1 error 50 | if rf, ok := ret.Get(1).(func(context.Context, *datacatalog.AddTagRequest) error); ok { 51 | r1 = rf(ctx, request) 52 | } else { 53 | r1 = ret.Error(1) 54 | } 55 | 56 | return r0, r1 57 | } 58 | -------------------------------------------------------------------------------- /pkg/repositories/config/postgres.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "github.com/flyteorg/flytestdlib/database" 8 | stdlibLogger "github.com/flyteorg/flytestdlib/logger" 9 | "github.com/flyteorg/flytestdlib/promutils" 10 | 11 | "gorm.io/driver/postgres" 12 | "gorm.io/gorm" 13 | "gorm.io/gorm/logger" 14 | ) 15 | 16 | const ( 17 | Postgres = "postgres" 18 | Sqlite = "sqlite" 19 | ) 20 | 21 | // Generic interface for providing a config necessary to open a database connection. 22 | type DbConnectionConfigProvider interface { 23 | // Returns database dialector 24 | GetDialector() gorm.Dialector 25 | 26 | GetDBConfig() database.DbConfig 27 | 28 | GetDSN() string 29 | } 30 | 31 | type BaseConfig struct { 32 | LogLevel logger.LogLevel `json:"log_level"` 33 | DisableForeignKeyConstraintWhenMigrating bool 34 | } 35 | 36 | // PostgreSQL implementation for DbConnectionConfigProvider. 37 | type PostgresConfigProvider struct { 38 | config database.DbConfig 39 | scope promutils.Scope 40 | } 41 | 42 | // TODO : Make the Config provider itself env based 43 | func NewPostgresConfigProvider(config database.DbConfig, scope promutils.Scope) DbConnectionConfigProvider { 44 | return &PostgresConfigProvider{ 45 | config: config, 46 | scope: scope, 47 | } 48 | } 49 | 50 | func (p *PostgresConfigProvider) GetDSN() string { 51 | if p.config.Postgres.Password == "" { 52 | // Switch for development 53 | return fmt.Sprintf("host=%s port=%d dbname=%s user=%s sslmode=disable", 54 | p.config.Postgres.Host, p.config.Postgres.Port, p.config.Postgres.DbName, p.config.Postgres.User) 55 | } 56 | return fmt.Sprintf("host=%s port=%d dbname=%s user=%s password=%s %s", 57 | p.config.Postgres.Host, p.config.Postgres.Port, p.config.Postgres.DbName, p.config.Postgres.User, p.config.Postgres.Password, p.config.Postgres.ExtraOptions) 58 | } 59 | 60 | func (p *PostgresConfigProvider) GetDialector() gorm.Dialector { 61 | return postgres.Open(p.GetDSN()) 62 | } 63 | 64 | func (p *PostgresConfigProvider) GetDBConfig() database.DbConfig { 65 | return p.config 66 | } 67 | 68 | // Opens a connection to the database specified in the config. 69 | // You must call CloseDbConnection at the end of your session! 70 | func OpenDbConnection(ctx context.Context, config DbConnectionConfigProvider) (*gorm.DB, error) { 71 | dbConfig := config.GetDBConfig() 72 | 73 | db, err := gorm.Open(config.GetDialector(), &gorm.Config{ 74 | Logger: database.GetGormLogger(ctx, stdlibLogger.GetConfig()), 75 | DisableForeignKeyConstraintWhenMigrating: !dbConfig.EnableForeignKeyConstraintWhenMigrating, 76 | }) 77 | if err != nil { 78 | return nil, err 79 | } 80 | 81 | return db, setupDbConnectionPool(db, &dbConfig) 82 | } 83 | 84 | func setupDbConnectionPool(gormDb *gorm.DB, dbConfig *database.DbConfig) error { 85 | genericDb, err := gormDb.DB() 86 | if err != nil { 87 | return err 88 | } 89 | genericDb.SetConnMaxLifetime(dbConfig.ConnMaxLifeTime.Duration) 90 | genericDb.SetMaxIdleConns(dbConfig.MaxIdleConnections) 91 | genericDb.SetMaxOpenConns(dbConfig.MaxOpenConnections) 92 | return nil 93 | } 94 | -------------------------------------------------------------------------------- /pkg/repositories/config/postgres_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "os" 5 | "path/filepath" 6 | "testing" 7 | "time" 8 | 9 | "github.com/flyteorg/flytestdlib/config" 10 | "github.com/flyteorg/flytestdlib/database" 11 | mockScope "github.com/flyteorg/flytestdlib/promutils" 12 | 13 | "github.com/stretchr/testify/assert" 14 | 15 | "gorm.io/driver/sqlite" 16 | "gorm.io/gorm" 17 | ) 18 | 19 | func TestConstructGormArgs(t *testing.T) { 20 | postgresConfigProvider := NewPostgresConfigProvider(database.DbConfig{Postgres: database.PostgresConfig{ 21 | Host: "localhost", 22 | Port: 5432, 23 | DbName: "postgres", 24 | User: "postgres", 25 | ExtraOptions: "sslmode=disable", 26 | }, 27 | EnableForeignKeyConstraintWhenMigrating: false, 28 | }, mockScope.NewTestScope()) 29 | 30 | assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres sslmode=disable", postgresConfigProvider.GetDSN()) 31 | assert.Equal(t, false, postgresConfigProvider.GetDBConfig().EnableForeignKeyConstraintWhenMigrating) 32 | } 33 | 34 | func TestConstructGormArgsWithPassword(t *testing.T) { 35 | postgresConfigProvider := NewPostgresConfigProvider(database.DbConfig{Postgres: database.PostgresConfig{ 36 | Host: "localhost", 37 | Port: 5432, 38 | DbName: "postgres", 39 | User: "postgres", 40 | Password: "pass", 41 | ExtraOptions: "sslmode=enable", 42 | }, 43 | }, mockScope.NewTestScope()) 44 | 45 | assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=pass sslmode=enable", postgresConfigProvider.GetDSN()) 46 | } 47 | 48 | func TestConstructGormArgsWithPasswordNoExtra(t *testing.T) { 49 | postgresConfigProvider := NewPostgresConfigProvider(database.DbConfig{Postgres: database.PostgresConfig{ 50 | Host: "localhost", 51 | Port: 5432, 52 | DbName: "postgres", 53 | User: "postgres", 54 | Password: "pass", 55 | }, 56 | }, mockScope.NewTestScope()) 57 | 58 | assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=pass ", postgresConfigProvider.GetDSN()) 59 | } 60 | 61 | func TestSetupDbConnectionPool(t *testing.T) { 62 | t.Run("successful", func(t *testing.T) { 63 | gormDb, err := gorm.Open(sqlite.Open(filepath.Join(os.TempDir(), "gorm.db")), &gorm.Config{}) 64 | assert.Nil(t, err) 65 | dbConfig := &database.DbConfig{ 66 | DeprecatedPort: 5432, 67 | MaxIdleConnections: 10, 68 | MaxOpenConnections: 1000, 69 | ConnMaxLifeTime: config.Duration{Duration: time.Hour}, 70 | } 71 | err = setupDbConnectionPool(gormDb, dbConfig) 72 | assert.Nil(t, err) 73 | genericDb, err := gormDb.DB() 74 | assert.Nil(t, err) 75 | assert.Equal(t, genericDb.Stats().MaxOpenConnections, 1000) 76 | }) 77 | t.Run("unset duration", func(t *testing.T) { 78 | gormDb, err := gorm.Open(sqlite.Open(filepath.Join(os.TempDir(), "gorm.db")), &gorm.Config{}) 79 | assert.Nil(t, err) 80 | dbConfig := &database.DbConfig{ 81 | DeprecatedPort: 5432, 82 | MaxIdleConnections: 10, 83 | MaxOpenConnections: 1000, 84 | } 85 | err = setupDbConnectionPool(gormDb, dbConfig) 86 | assert.Nil(t, err) 87 | genericDb, err := gormDb.DB() 88 | assert.Nil(t, err) 89 | assert.Equal(t, genericDb.Stats().MaxOpenConnections, 1000) 90 | }) 91 | t.Run("failed to get DB", func(t *testing.T) { 92 | gormDb := &gorm.DB{ 93 | Config: &gorm.Config{ 94 | ConnPool: &gorm.PreparedStmtDB{}, 95 | }, 96 | } 97 | dbConfig := &database.DbConfig{ 98 | DeprecatedPort: 5432, 99 | MaxIdleConnections: 10, 100 | MaxOpenConnections: 1000, 101 | ConnMaxLifeTime: config.Duration{Duration: time.Hour}, 102 | } 103 | err := setupDbConnectionPool(gormDb, dbConfig) 104 | assert.NotNil(t, err) 105 | }) 106 | } 107 | -------------------------------------------------------------------------------- /pkg/repositories/errors/error_transformer.go: -------------------------------------------------------------------------------- 1 | package errors 2 | 3 | // Defines the basic error transformer interface that all database types must implement. 4 | type ErrorTransformer interface { 5 | ToDataCatalogError(err error) error 6 | } 7 | -------------------------------------------------------------------------------- /pkg/repositories/errors/errors.go: -------------------------------------------------------------------------------- 1 | // Generic errors used in the repos layer 2 | package errors 3 | 4 | import ( 5 | "github.com/flyteorg/datacatalog/pkg/common" 6 | "github.com/flyteorg/datacatalog/pkg/errors" 7 | "github.com/golang/protobuf/proto" 8 | "google.golang.org/grpc/codes" 9 | ) 10 | 11 | const ( 12 | AlreadyExists = "entity already exists" 13 | notFound = "missing entity of type %s with identifier %v" 14 | invalidJoin = "cannot relate entity %s with entity %s" 15 | invalidEntity = "no such entity %s" 16 | ) 17 | 18 | func GetMissingEntityError(entityType string, identifier proto.Message) error { 19 | return errors.NewDataCatalogErrorf(codes.NotFound, notFound, entityType, identifier) 20 | } 21 | 22 | func GetInvalidEntityRelationshipError(entityType common.Entity, otherEntityType common.Entity) error { 23 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, invalidJoin, entityType, otherEntityType) 24 | } 25 | 26 | func GetInvalidEntityError(entityType common.Entity) error { 27 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, invalidEntity, entityType) 28 | } 29 | 30 | func GetUnsupportedFilterExpressionErr(operator common.ComparisonOperator) error { 31 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, "unsupported filter expression operator index: %v", 32 | operator) 33 | } 34 | -------------------------------------------------------------------------------- /pkg/repositories/errors/postgres.go: -------------------------------------------------------------------------------- 1 | package errors 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "reflect" 7 | 8 | "github.com/flyteorg/flytestdlib/logger" 9 | 10 | "github.com/jackc/pgconn" 11 | 12 | catalogErrors "github.com/flyteorg/datacatalog/pkg/errors" 13 | "google.golang.org/grpc/codes" 14 | "gorm.io/gorm" 15 | ) 16 | 17 | // Postgres error codes 18 | const ( 19 | uniqueConstraintViolationCode = "23505" 20 | undefinedTable = "42P01" 21 | ) 22 | 23 | type postgresErrorTransformer struct { 24 | } 25 | 26 | const ( 27 | unexpectedType = "unexpected error type for: %v" 28 | uniqueConstraintViolation = "value with matching already exists (%s)" 29 | defaultPgError = "failed database operation with code [%s] and msg [%s]" 30 | unsupportedTableOperation = "cannot query with specified table attributes: %s" 31 | ) 32 | 33 | func (p *postgresErrorTransformer) fromGormError(err error) error { 34 | switch err.Error() { 35 | case gorm.ErrRecordNotFound.Error(): 36 | return catalogErrors.NewDataCatalogErrorf(codes.NotFound, "entry not found") 37 | default: 38 | return catalogErrors.NewDataCatalogErrorf(codes.Internal, unexpectedType, err) 39 | } 40 | } 41 | 42 | func (p *postgresErrorTransformer) ToDataCatalogError(err error) error { 43 | if unwrappedErr := errors.Unwrap(err); unwrappedErr != nil { 44 | err = unwrappedErr 45 | } 46 | 47 | pqError, ok := err.(*pgconn.PgError) 48 | if !ok { 49 | logger.InfofNoCtx("Unable to cast to pgconn.PgError. Error type: [%v]", 50 | reflect.TypeOf(err)) 51 | return p.fromGormError(err) 52 | } 53 | 54 | switch pqError.Code { 55 | case uniqueConstraintViolationCode: 56 | return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, uniqueConstraintViolation, pqError.Message) 57 | case undefinedTable: 58 | return catalogErrors.NewDataCatalogErrorf(codes.InvalidArgument, unsupportedTableOperation, pqError.Message) 59 | default: 60 | return catalogErrors.NewDataCatalogErrorf(codes.Unknown, fmt.Sprintf(defaultPgError, pqError.Code, pqError.Message)) 61 | } 62 | } 63 | 64 | func NewPostgresErrorTransformer() ErrorTransformer { 65 | return &postgresErrorTransformer{} 66 | } 67 | 68 | type ConnectError interface { 69 | Unwrap() error 70 | Error() string 71 | } 72 | -------------------------------------------------------------------------------- /pkg/repositories/factory.go: -------------------------------------------------------------------------------- 1 | package repositories 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "github.com/flyteorg/flytestdlib/database" 8 | 9 | "github.com/flyteorg/datacatalog/pkg/repositories/config" 10 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 11 | "github.com/flyteorg/datacatalog/pkg/repositories/interfaces" 12 | "github.com/flyteorg/flytestdlib/promutils" 13 | ) 14 | 15 | type RepoConfig int32 16 | 17 | const ( 18 | POSTGRES RepoConfig = 0 19 | ) 20 | 21 | var RepositoryConfigurationName = map[RepoConfig]string{ 22 | POSTGRES: "POSTGRES", 23 | } 24 | 25 | // The RepositoryInterface indicates the methods that each Repository must support. 26 | // A Repository indicates a Database which is collection of Tables/models. 27 | // The goal is allow databases to be Plugged in easily. 28 | type RepositoryInterface interface { 29 | DatasetRepo() interfaces.DatasetRepo 30 | ArtifactRepo() interfaces.ArtifactRepo 31 | TagRepo() interfaces.TagRepo 32 | ReservationRepo() interfaces.ReservationRepo 33 | } 34 | 35 | func GetRepository(ctx context.Context, repoType RepoConfig, dbConfig database.DbConfig, scope promutils.Scope) RepositoryInterface { 36 | switch repoType { 37 | case POSTGRES: 38 | db, err := config.OpenDbConnection(ctx, config.NewPostgresConfigProvider(dbConfig, scope.NewSubScope("postgres"))) 39 | if err != nil { 40 | panic(err) 41 | } 42 | return NewPostgresRepo( 43 | db, 44 | errors.NewPostgresErrorTransformer(), 45 | scope.NewSubScope("repositories")) 46 | default: 47 | panic(fmt.Sprintf("Invalid repoType %v", repoType)) 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/dataset.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "context" 5 | 6 | idl_datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | 8 | "github.com/flyteorg/datacatalog/pkg/common" 9 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 10 | "github.com/flyteorg/datacatalog/pkg/repositories/interfaces" 11 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 12 | "github.com/flyteorg/flytestdlib/logger" 13 | "github.com/flyteorg/flytestdlib/promutils" 14 | "gorm.io/gorm" 15 | ) 16 | 17 | type dataSetRepo struct { 18 | db *gorm.DB 19 | errorTransformer errors.ErrorTransformer 20 | repoMetrics gormMetrics 21 | } 22 | 23 | func NewDatasetRepo(db *gorm.DB, errorTransformer errors.ErrorTransformer, scope promutils.Scope) interfaces.DatasetRepo { 24 | return &dataSetRepo{ 25 | db: db, 26 | errorTransformer: errorTransformer, 27 | repoMetrics: newGormMetrics(scope), 28 | } 29 | } 30 | 31 | // Create a Dataset model 32 | func (h *dataSetRepo) Create(ctx context.Context, in models.Dataset) error { 33 | timer := h.repoMetrics.CreateDuration.Start(ctx) 34 | defer timer.Stop() 35 | 36 | result := h.db.Create(&in) 37 | if result.Error != nil { 38 | return h.errorTransformer.ToDataCatalogError(result.Error) 39 | } 40 | return nil 41 | } 42 | 43 | // Get Dataset model 44 | func (h *dataSetRepo) Get(ctx context.Context, in models.DatasetKey) (models.Dataset, error) { 45 | timer := h.repoMetrics.GetDuration.Start(ctx) 46 | defer timer.Stop() 47 | 48 | var ds models.Dataset 49 | result := h.db.Preload("PartitionKeys", func(db *gorm.DB) *gorm.DB { 50 | return db.Order("partition_keys.created_at ASC") // preserve the order in which the partitions were created 51 | }).First(&ds, &models.Dataset{DatasetKey: in}) 52 | 53 | if result.Error != nil { 54 | logger.Debugf(ctx, "Unable to find Dataset: [%+v], err: %v", in, result.Error) 55 | 56 | if result.Error.Error() == gorm.ErrRecordNotFound.Error() { 57 | return models.Dataset{}, errors.GetMissingEntityError("Dataset", &idl_datacatalog.DatasetID{ 58 | Project: in.Project, 59 | Domain: in.Domain, 60 | Name: in.Name, 61 | Version: in.Version, 62 | }) 63 | } 64 | return models.Dataset{}, h.errorTransformer.ToDataCatalogError(result.Error) 65 | } 66 | 67 | return ds, nil 68 | } 69 | 70 | func (h *dataSetRepo) List(ctx context.Context, in models.ListModelsInput) ([]models.Dataset, error) { 71 | timer := h.repoMetrics.ListDuration.Start(ctx) 72 | defer timer.Stop() 73 | 74 | // apply filters and joins 75 | tx, err := applyListModelsInput(h.db, common.Dataset, in) 76 | if err != nil { 77 | return nil, err 78 | } else if tx.Error != nil { 79 | return []models.Dataset{}, h.errorTransformer.ToDataCatalogError(tx.Error) 80 | } 81 | 82 | datasets := make([]models.Dataset, 0) 83 | tx = tx.Preload("PartitionKeys").Find(&datasets) 84 | if tx.Error != nil { 85 | return []models.Dataset{}, h.errorTransformer.ToDataCatalogError(tx.Error) 86 | } 87 | return datasets, nil 88 | } 89 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/filter.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/common" 7 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 8 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 9 | ) 10 | 11 | // String formats for various GORM expression queries 12 | const ( 13 | equalQuery = "%s.%s = ?" 14 | ) 15 | 16 | type gormValueFilterImpl struct { 17 | comparisonOperator common.ComparisonOperator 18 | field string 19 | value interface{} 20 | } 21 | 22 | // Get the GORM expression to filter by a model's property. The output should be a valid input into tx.Where() 23 | func (g *gormValueFilterImpl) GetDBQueryExpression(tableName string) (models.DBQueryExpr, error) { 24 | switch g.comparisonOperator { 25 | case common.Equal: 26 | return models.DBQueryExpr{ 27 | Query: fmt.Sprintf(equalQuery, tableName, g.field), 28 | Args: g.value, 29 | }, nil 30 | } 31 | return models.DBQueryExpr{}, errors.GetUnsupportedFilterExpressionErr(g.comparisonOperator) 32 | } 33 | 34 | // Construct the container necessary to issue a db query to filter in GORM 35 | func NewGormValueFilter(comparisonOperator common.ComparisonOperator, field string, value interface{}) models.ModelValueFilter { 36 | return &gormValueFilterImpl{ 37 | comparisonOperator: comparisonOperator, 38 | field: field, 39 | value: value, 40 | } 41 | } 42 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/filter_test.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/common" 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestGormValueFilter(t *testing.T) { 11 | filter := NewGormValueFilter(common.Equal, "key", "region") 12 | expression, err := filter.GetDBQueryExpression("partitions") 13 | assert.NoError(t, err) 14 | assert.Equal(t, expression.Query, "partitions.key = ?") 15 | assert.Equal(t, expression.Args, "region") 16 | } 17 | 18 | func TestGormValueFilterInvalidOperator(t *testing.T) { 19 | filter := NewGormValueFilter(123, "key", "region") 20 | _, err := filter.GetDBQueryExpression("partitions") 21 | assert.Error(t, err) 22 | } 23 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/join.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "fmt" 5 | "strings" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/common" 8 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 9 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 10 | ) 11 | 12 | const ( 13 | joinCondition = "JOIN %s %s ON %s" // Format for the join: JOIN
ON 14 | joinEquals = "%s.%s = %s.%s" // Format for the columns to join: . = . 15 | joinSeparator = " AND " // Separator if there's more than one joining column 16 | ) 17 | 18 | // JoinOnMap is a map of the properties for joining source table to joining table 19 | type JoinOnMap map[string]string 20 | 21 | // This provides the field names needed for joining a source Model to joining Model 22 | var joinFieldNames = map[common.Entity]map[common.Entity]JoinOnMap{ 23 | common.Artifact: { 24 | common.Partition: JoinOnMap{"artifact_id": "artifact_id"}, 25 | common.Tag: JoinOnMap{"artifact_id": "artifact_id"}, 26 | }, 27 | } 28 | 29 | // Contains the details to construct GORM JOINs in the format: 30 | // JOIN sourceTable ON sourceTable.sourceField = joiningTable.joiningField 31 | type gormJoinConditionImpl struct { 32 | // The source entity type 33 | sourceEntity common.Entity 34 | // The joining entity type 35 | joiningEntity common.Entity 36 | } 37 | 38 | // Get the GORM expression to JOIN two entities. The output should be a valid input into tx.Join() 39 | func (g *gormJoinConditionImpl) GetJoinOnDBQueryExpression(sourceTableName string, joiningTableName string, joiningTableAlias string) (string, error) { 40 | joinOnFieldMap, err := g.getJoinOnFields() 41 | 42 | if err != nil { 43 | return "", err 44 | } 45 | 46 | joinFields := make([]string, 0, len(joinOnFieldMap)) 47 | for sourceField, joiningField := range joinOnFieldMap { 48 | joinFieldCondition := fmt.Sprintf(joinEquals, sourceTableName, sourceField, joiningTableAlias, joiningField) 49 | joinFields = append(joinFields, joinFieldCondition) 50 | } 51 | 52 | return fmt.Sprintf(joinCondition, joiningTableName, joiningTableAlias, strings.Join(joinFields, joinSeparator)), nil 53 | } 54 | 55 | // Get the properties necessary to join two GORM models 56 | func (g *gormJoinConditionImpl) getJoinOnFields() (JoinOnMap, error) { 57 | joiningEntityMap, ok := joinFieldNames[g.sourceEntity] 58 | if !ok { 59 | return nil, errors.GetInvalidEntityRelationshipError(g.sourceEntity, g.joiningEntity) 60 | } 61 | 62 | fieldMap, ok := joiningEntityMap[g.joiningEntity] 63 | if !ok { 64 | return nil, errors.GetInvalidEntityRelationshipError(g.sourceEntity, g.joiningEntity) 65 | } 66 | 67 | return fieldMap, nil 68 | } 69 | 70 | func NewGormJoinCondition(sourceEntity common.Entity, joiningEntity common.Entity) models.ModelJoinCondition { 71 | return &gormJoinConditionImpl{ 72 | joiningEntity: joiningEntity, 73 | sourceEntity: sourceEntity, 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/join_test.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/common" 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestGormJoinCondition(t *testing.T) { 11 | filter := NewGormJoinCondition(common.Artifact, common.Partition) 12 | joinQuery, err := filter.GetJoinOnDBQueryExpression("artifacts", "partitions", "p") 13 | assert.NoError(t, err) 14 | assert.Equal(t, joinQuery, "JOIN partitions p ON artifacts.artifact_id = p.artifact_id") 15 | } 16 | 17 | // Tag cannot be joined with partitions 18 | func TestInvalidGormJoinCondition(t *testing.T) { 19 | filter := NewGormJoinCondition(common.Tag, common.Partition) 20 | 21 | _, err := filter.GetJoinOnDBQueryExpression("tags", "partitions", "t") 22 | assert.Error(t, err) 23 | } 24 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/list.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "fmt" 5 | 6 | errors2 "github.com/flyteorg/datacatalog/pkg/errors" 7 | 8 | "github.com/flyteorg/datacatalog/pkg/common" 9 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 10 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | "google.golang.org/grpc/codes" 12 | "gorm.io/gorm" 13 | ) 14 | 15 | const ( 16 | tableAliasFormat = "%s%d" // Table Alias is the "
" 17 | ) 18 | 19 | var entityToModel = map[common.Entity]interface{}{ 20 | common.Artifact: models.Artifact{}, 21 | common.Dataset: models.Dataset{}, 22 | common.Partition: models.Partition{}, 23 | common.Tag: models.Tag{}, 24 | } 25 | 26 | func getTableName(tx *gorm.DB, model interface{}) (string, error) { 27 | stmt := gorm.Statement{DB: tx} 28 | 29 | if err := stmt.Parse(model); err != nil { 30 | return "", errors2.NewDataCatalogError(codes.InvalidArgument, err.Error()) 31 | } 32 | return stmt.Schema.Table, nil 33 | } 34 | 35 | // Apply the list query on the source model. This method will apply the necessary joins, filters and 36 | // pagination on the database for the given ListModelInputs. 37 | func applyListModelsInput(tx *gorm.DB, sourceEntity common.Entity, in models.ListModelsInput) (*gorm.DB, error) { 38 | sourceModel, ok := entityToModel[sourceEntity] 39 | if !ok { 40 | return nil, errors.GetInvalidEntityError(sourceEntity) 41 | } 42 | 43 | sourceTableName, err := getTableName(tx, sourceModel) 44 | if err != nil { 45 | return nil, err 46 | } 47 | 48 | for modelIndex, modelFilter := range in.ModelFilters { 49 | entity := modelFilter.Entity 50 | filterModel, ok := entityToModel[entity] 51 | if !ok { 52 | return nil, errors.GetInvalidEntityError(entity) 53 | } 54 | tableName, err := getTableName(tx, filterModel) 55 | if err != nil { 56 | return nil, err 57 | } 58 | tableAlias := tableName 59 | 60 | // Optionally add the join condition if the entity we need isn't the source 61 | if sourceEntity != modelFilter.Entity { 62 | // if there is a join associated with the filter, we should use an alias 63 | joinCondition := modelFilter.JoinCondition 64 | tableAlias = fmt.Sprintf(tableAliasFormat, tableName, modelIndex) 65 | joinExpression, err := joinCondition.GetJoinOnDBQueryExpression(sourceTableName, tableName, tableAlias) 66 | if err != nil { 67 | return nil, err 68 | } 69 | tx = tx.Joins(joinExpression) 70 | } 71 | 72 | for _, whereFilter := range modelFilter.ValueFilters { 73 | dbQueryExpr, err := whereFilter.GetDBQueryExpression(tableAlias) 74 | 75 | if err != nil { 76 | return nil, err 77 | } 78 | tx = tx.Where(dbQueryExpr.Query, dbQueryExpr.Args) 79 | } 80 | } 81 | 82 | tx = tx.Limit(in.Limit) 83 | tx = tx.Offset(in.Offset) 84 | 85 | if in.SortParameter != nil { 86 | tx = tx.Order(in.SortParameter.GetDBOrderExpression(sourceTableName)) 87 | } 88 | return tx, nil 89 | } 90 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/list_test.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "database/sql/driver" 5 | "strings" 6 | "testing" 7 | 8 | mocket "github.com/Selvatico/go-mocket" 9 | "github.com/flyteorg/datacatalog/pkg/common" 10 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | "github.com/flyteorg/datacatalog/pkg/repositories/utils" 12 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 13 | "github.com/stretchr/testify/assert" 14 | ) 15 | 16 | func TestApplyFilter(t *testing.T) { 17 | testDB := utils.GetDbForTest(t) 18 | GlobalMock := mocket.Catcher.Reset() 19 | GlobalMock.Logging = true 20 | validInputApply := false 21 | 22 | GlobalMock.NewMock().WithQuery( 23 | `SELECT "artifacts"."created_at","artifacts"."updated_at","artifacts"."deleted_at","artifacts"."dataset_project","artifacts"."dataset_name","artifacts"."dataset_domain","artifacts"."dataset_version","artifacts"."artifact_id","artifacts"."dataset_uuid","artifacts"."serialized_metadata" FROM "artifacts"`).WithCallback( 24 | func(s string, values []driver.NamedValue) { 25 | // separate the regex matching because the joins reorder on different test runs 26 | validInputApply = strings.Contains(s, `JOIN tags tags1 ON artifacts.artifact_id = tags1.artifact_id`) && 27 | strings.Contains(s, `JOIN partitions partitions0 ON artifacts.artifact_id = partitions0.artifact_id`) && 28 | strings.Contains(s, `WHERE partitions0.key1 = $1 AND partitions0.key2 = $2 AND tags1.tag_name = $3 `+ 29 | `ORDER BY artifacts.created_at desc LIMIT 10 OFFSET 10`) 30 | }) 31 | 32 | listInput := models.ListModelsInput{ 33 | ModelFilters: []models.ModelFilter{ 34 | { 35 | Entity: common.Partition, 36 | JoinCondition: NewGormJoinCondition(common.Artifact, common.Partition), 37 | ValueFilters: []models.ModelValueFilter{ 38 | NewGormValueFilter(common.Equal, "key1", "val1"), 39 | NewGormValueFilter(common.Equal, "key2", "val2"), 40 | }, 41 | }, 42 | { 43 | Entity: common.Tag, 44 | JoinCondition: NewGormJoinCondition(common.Artifact, common.Tag), 45 | ValueFilters: []models.ModelValueFilter{ 46 | NewGormValueFilter(common.Equal, "tag_name", "special"), 47 | }, 48 | }, 49 | }, 50 | Offset: 10, 51 | Limit: 10, 52 | SortParameter: NewGormSortParameter(datacatalog.PaginationOptions_CREATION_TIME, datacatalog.PaginationOptions_DESCENDING), 53 | } 54 | 55 | tx, err := applyListModelsInput(testDB, common.Artifact, listInput) 56 | assert.NoError(t, err) 57 | 58 | tx.Find(models.Artifact{}) 59 | assert.True(t, validInputApply) 60 | } 61 | 62 | func TestApplyFilterEmpty(t *testing.T) { 63 | testDB := utils.GetDbForTest(t) 64 | GlobalMock := mocket.Catcher.Reset() 65 | GlobalMock.Logging = true 66 | validInputApply := false 67 | 68 | GlobalMock.NewMock().WithQuery( 69 | `SELECT * FROM "artifacts" LIMIT 10 OFFSET 10`).WithCallback( 70 | func(s string, values []driver.NamedValue) { 71 | // separate the regex matching because the joins reorder on different test runs 72 | validInputApply = true 73 | }) 74 | 75 | listInput := models.ListModelsInput{ 76 | Offset: 10, 77 | Limit: 10, 78 | } 79 | 80 | tx, err := applyListModelsInput(testDB, common.Artifact, listInput) 81 | assert.NoError(t, err) 82 | 83 | tx.Find(models.Artifact{}) 84 | assert.True(t, validInputApply) 85 | } 86 | 87 | func TestGetTableErr(t *testing.T) { 88 | testDB := utils.GetDbForTest(t) 89 | 90 | tableName, err := getTableName(testDB, "") 91 | assert.Error(t, err) 92 | assert.Equal(t, "", tableName) 93 | } 94 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/metrics.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "time" 5 | 6 | "github.com/flyteorg/flytestdlib/promutils" 7 | "github.com/flyteorg/flytestdlib/promutils/labeled" 8 | ) 9 | 10 | // Common metrics for DB CRUD operations 11 | type gormMetrics struct { 12 | Scope promutils.Scope 13 | CreateDuration labeled.StopWatch 14 | DeleteDuration labeled.StopWatch 15 | GetDuration labeled.StopWatch 16 | ListDuration labeled.StopWatch 17 | UpdateDuration labeled.StopWatch 18 | } 19 | 20 | func newGormMetrics(scope promutils.Scope) gormMetrics { 21 | return gormMetrics{ 22 | Scope: scope, 23 | CreateDuration: labeled.NewStopWatch( 24 | "create", "Duration for creating a new entity", time.Millisecond, scope), 25 | DeleteDuration: labeled.NewStopWatch( 26 | "delete", "Duration for deleting a new entity", time.Millisecond, scope), 27 | GetDuration: labeled.NewStopWatch( 28 | "get", "Duration for retrieving an entity ", time.Millisecond, scope), 29 | ListDuration: labeled.NewStopWatch( 30 | "list", "Duration for listing entities ", time.Millisecond, scope), 31 | UpdateDuration: labeled.NewStopWatch( 32 | "update", "Duration for updating entities ", time.Millisecond, scope), 33 | } 34 | } 35 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/reservation.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "context" 5 | 6 | datacatalog_error "github.com/flyteorg/datacatalog/pkg/errors" 7 | "google.golang.org/grpc/codes" 8 | 9 | "time" 10 | 11 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 12 | 13 | errors2 "github.com/flyteorg/datacatalog/pkg/repositories/errors" 14 | "github.com/flyteorg/datacatalog/pkg/repositories/interfaces" 15 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 16 | "github.com/flyteorg/flytestdlib/promutils" 17 | "gorm.io/gorm" 18 | "gorm.io/gorm/clause" 19 | ) 20 | 21 | type reservationRepo struct { 22 | db *gorm.DB 23 | repoMetrics gormMetrics 24 | errorTransformer errors2.ErrorTransformer 25 | } 26 | 27 | // NewReservationRepo creates a reservationRepo 28 | func NewReservationRepo(db *gorm.DB, errorTransformer errors2.ErrorTransformer, scope promutils.Scope) interfaces.ReservationRepo { 29 | return &reservationRepo{ 30 | db: db, 31 | errorTransformer: errorTransformer, 32 | repoMetrics: newGormMetrics(scope), 33 | } 34 | } 35 | 36 | func (r *reservationRepo) Create(ctx context.Context, reservation models.Reservation, now time.Time) error { 37 | timer := r.repoMetrics.CreateDuration.Start(ctx) 38 | defer timer.Stop() 39 | 40 | result := r.db.Clauses(clause.OnConflict{DoNothing: true}).Create(&reservation) 41 | if result.Error != nil { 42 | return r.errorTransformer.ToDataCatalogError(result.Error) 43 | } 44 | 45 | if result.RowsAffected == 0 { 46 | return datacatalog_error.NewDataCatalogError(codes.FailedPrecondition, errors2.AlreadyExists) 47 | } 48 | 49 | return nil 50 | } 51 | 52 | func (r *reservationRepo) Delete(ctx context.Context, reservationKey models.ReservationKey, ownerID string) error { 53 | timer := r.repoMetrics.DeleteDuration.Start(ctx) 54 | defer timer.Stop() 55 | 56 | var reservation models.Reservation 57 | 58 | result := r.db.Where(&models.Reservation{ 59 | ReservationKey: reservationKey, 60 | OwnerID: ownerID, 61 | }).Delete(&reservation) 62 | if result.Error != nil { 63 | return r.errorTransformer.ToDataCatalogError(result.Error) 64 | } 65 | 66 | if result.RowsAffected == 0 { 67 | return errors2.GetMissingEntityError("Reservation", 68 | &datacatalog.ReservationID{ 69 | DatasetId: &datacatalog.DatasetID{ 70 | Project: reservationKey.DatasetProject, 71 | Domain: reservationKey.DatasetDomain, 72 | Name: reservationKey.DatasetName, 73 | Version: reservationKey.DatasetVersion, 74 | }, 75 | TagName: reservationKey.TagName, 76 | }) 77 | } 78 | 79 | return nil 80 | } 81 | 82 | func (r *reservationRepo) Get(ctx context.Context, reservationKey models.ReservationKey) (models.Reservation, error) { 83 | timer := r.repoMetrics.GetDuration.Start(ctx) 84 | defer timer.Stop() 85 | 86 | var reservation models.Reservation 87 | 88 | result := r.db.Where(&models.Reservation{ 89 | ReservationKey: reservationKey, 90 | }).Take(&reservation) 91 | 92 | if result.Error != nil { 93 | return reservation, r.errorTransformer.ToDataCatalogError(result.Error) 94 | } 95 | 96 | return reservation, nil 97 | } 98 | 99 | func (r *reservationRepo) Update(ctx context.Context, reservation models.Reservation, now time.Time) error { 100 | timer := r.repoMetrics.UpdateDuration.Start(ctx) 101 | defer timer.Stop() 102 | 103 | result := r.db.Model(&models.Reservation{ 104 | ReservationKey: reservation.ReservationKey, 105 | }).Where("expires_at<=? OR owner_id=?", now, reservation.OwnerID).Updates(reservation) 106 | if result.Error != nil { 107 | return r.errorTransformer.ToDataCatalogError(result.Error) 108 | } 109 | 110 | if result.RowsAffected == 0 { 111 | return datacatalog_error.NewDataCatalogError(codes.FailedPrecondition, errors2.AlreadyExists) 112 | } 113 | 114 | return nil 115 | } 116 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/sort.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 8 | ) 9 | 10 | const ( 11 | sortQuery = "%s.%s %s" 12 | ) 13 | 14 | // Container for the sort details 15 | type sortParameter struct { 16 | sortKey datacatalog.PaginationOptions_SortKey 17 | sortOrder datacatalog.PaginationOptions_SortOrder 18 | } 19 | 20 | // Generate the DBOrderExpression that GORM needs to order models 21 | func (s *sortParameter) GetDBOrderExpression(tableName string) string { 22 | var sortOrderString string 23 | switch s.sortOrder { 24 | case datacatalog.PaginationOptions_ASCENDING: 25 | sortOrderString = "asc" 26 | case datacatalog.PaginationOptions_DESCENDING: 27 | fallthrough 28 | default: 29 | sortOrderString = "desc" 30 | } 31 | 32 | var sortKeyString string 33 | switch s.sortKey { 34 | case datacatalog.PaginationOptions_CREATION_TIME: 35 | fallthrough 36 | default: 37 | sortKeyString = "created_at" 38 | } 39 | return fmt.Sprintf(sortQuery, tableName, sortKeyString, sortOrderString) 40 | } 41 | 42 | // Create SortParameter for GORM 43 | func NewGormSortParameter(sortKey datacatalog.PaginationOptions_SortKey, sortOrder datacatalog.PaginationOptions_SortOrder) models.SortParameter { 44 | return &sortParameter{sortKey: sortKey, sortOrder: sortOrder} 45 | } 46 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/sort_test.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestSortAsc(t *testing.T) { 11 | dbSortExpression := NewGormSortParameter( 12 | datacatalog.PaginationOptions_CREATION_TIME, 13 | datacatalog.PaginationOptions_ASCENDING).GetDBOrderExpression("artifacts") 14 | 15 | assert.Equal(t, dbSortExpression, "artifacts.created_at asc") 16 | } 17 | 18 | func TestSortDesc(t *testing.T) { 19 | dbSortExpression := NewGormSortParameter( 20 | datacatalog.PaginationOptions_CREATION_TIME, 21 | datacatalog.PaginationOptions_DESCENDING).GetDBOrderExpression("artifacts") 22 | 23 | assert.Equal(t, dbSortExpression, "artifacts.created_at desc") 24 | } 25 | -------------------------------------------------------------------------------- /pkg/repositories/gormimpl/tag.go: -------------------------------------------------------------------------------- 1 | package gormimpl 2 | 3 | import ( 4 | "context" 5 | 6 | idl_datacatalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | 8 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 9 | "github.com/flyteorg/datacatalog/pkg/repositories/interfaces" 10 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | "github.com/flyteorg/flytestdlib/promutils" 12 | "gorm.io/gorm" 13 | ) 14 | 15 | type tagRepo struct { 16 | db *gorm.DB 17 | errorTransformer errors.ErrorTransformer 18 | repoMetrics gormMetrics 19 | } 20 | 21 | func NewTagRepo(db *gorm.DB, errorTransformer errors.ErrorTransformer, scope promutils.Scope) interfaces.TagRepo { 22 | return &tagRepo{ 23 | db: db, 24 | errorTransformer: errorTransformer, 25 | repoMetrics: newGormMetrics(scope), 26 | } 27 | } 28 | 29 | func (h *tagRepo) Create(ctx context.Context, tag models.Tag) error { 30 | timer := h.repoMetrics.CreateDuration.Start(ctx) 31 | defer timer.Stop() 32 | 33 | db := h.db.Create(&tag) 34 | 35 | if db.Error != nil { 36 | return h.errorTransformer.ToDataCatalogError(db.Error) 37 | } 38 | return nil 39 | } 40 | 41 | func (h *tagRepo) Get(ctx context.Context, in models.TagKey) (models.Tag, error) { 42 | timer := h.repoMetrics.GetDuration.Start(ctx) 43 | defer timer.Stop() 44 | 45 | var tag models.Tag 46 | result := h.db.Preload("Artifact"). 47 | Preload("Artifact.ArtifactData"). 48 | Preload("Artifact.Partitions", func(db *gorm.DB) *gorm.DB { 49 | return db.Order("partitions.created_at ASC") // preserve the order in which the partitions were created 50 | }). 51 | Preload("Artifact.Tags"). 52 | Order("tags.created_at DESC"). 53 | First(&tag, &models.Tag{ 54 | TagKey: in, 55 | }) 56 | 57 | if result.Error != nil { 58 | if result.Error.Error() == gorm.ErrRecordNotFound.Error() { 59 | return models.Tag{}, errors.GetMissingEntityError("Tag", &idl_datacatalog.Tag{ 60 | Name: tag.TagName, 61 | }) 62 | } 63 | return models.Tag{}, h.errorTransformer.ToDataCatalogError(result.Error) 64 | } 65 | 66 | return tag, nil 67 | } 68 | -------------------------------------------------------------------------------- /pkg/repositories/handle.go: -------------------------------------------------------------------------------- 1 | package repositories 2 | 3 | import ( 4 | "context" 5 | 6 | "gorm.io/driver/sqlite" 7 | 8 | "fmt" 9 | 10 | "github.com/flyteorg/flytestdlib/database" 11 | 12 | "github.com/flyteorg/datacatalog/pkg/repositories/config" 13 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 14 | "github.com/flyteorg/flytestdlib/logger" 15 | "github.com/flyteorg/flytestdlib/promutils" 16 | "gorm.io/gorm" 17 | ) 18 | 19 | type DBHandle struct { 20 | db *gorm.DB 21 | } 22 | 23 | func NewDBHandle(ctx context.Context, dbConfigValues database.DbConfig, catalogScope promutils.Scope) (*DBHandle, error) { 24 | var gormDb *gorm.DB 25 | var err error 26 | 27 | switch { 28 | case !dbConfigValues.SQLite.IsEmpty(): 29 | gormDb, err = gorm.Open(sqlite.Open(dbConfigValues.SQLite.File)) 30 | case !dbConfigValues.Postgres.IsEmpty(): 31 | gormDb, err = config.OpenDbConnection(ctx, config.NewPostgresConfigProvider(dbConfigValues, catalogScope.NewSubScope(config.Postgres))) 32 | default: 33 | return nil, fmt.Errorf("unrecognized database config, %v. Supported only postgres and sqlite", dbConfigValues) 34 | } 35 | 36 | if err != nil { 37 | return nil, err 38 | } 39 | 40 | out := &DBHandle{ 41 | db: gormDb, 42 | } 43 | 44 | return out, nil 45 | } 46 | 47 | func (h *DBHandle) CreateDB(dbName string) error { 48 | type DatabaseResult struct { 49 | Exists bool 50 | } 51 | var checkExists DatabaseResult 52 | result := h.db.Raw("SELECT EXISTS(SELECT datname FROM pg_catalog.pg_database WHERE datname = ?)", dbName).Scan(&checkExists) 53 | if result.Error != nil { 54 | return result.Error 55 | } 56 | 57 | // create db if it does not exist 58 | if !checkExists.Exists { 59 | logger.Infof(context.TODO(), "Creating Database %v since it does not exist", dbName) 60 | 61 | // NOTE: golang sql drivers do not support parameter injection for CREATE calls 62 | createDBStatement := fmt.Sprintf("CREATE DATABASE %s", dbName) 63 | result = h.db.Exec(createDBStatement) 64 | 65 | if result.Error != nil { 66 | if !isPgErrorWithCode(result.Error, pqDbAlreadyExistsCode) { 67 | return result.Error 68 | } 69 | logger.Infof(context.TODO(), "Not creating database %s, already exists", dbName) 70 | } 71 | } 72 | 73 | return nil 74 | } 75 | 76 | func (h *DBHandle) Migrate(ctx context.Context) error { 77 | if err := h.db.AutoMigrate(&models.Dataset{}); err != nil { 78 | return err 79 | } 80 | 81 | if err := h.db.Debug().AutoMigrate(&models.Artifact{}); err != nil { 82 | return err 83 | } 84 | 85 | if err := h.db.AutoMigrate(&models.ArtifactData{}); err != nil { 86 | return err 87 | } 88 | 89 | if err := h.db.AutoMigrate(&models.Tag{}); err != nil { 90 | return err 91 | } 92 | 93 | if err := h.db.AutoMigrate(&models.PartitionKey{}); err != nil { 94 | return err 95 | } 96 | 97 | if err := h.db.AutoMigrate(&models.Partition{}); err != nil { 98 | return err 99 | } 100 | 101 | if err := h.db.AutoMigrate(&models.Reservation{}); err != nil { 102 | return err 103 | } 104 | 105 | return nil 106 | } 107 | -------------------------------------------------------------------------------- /pkg/repositories/handle_test.go: -------------------------------------------------------------------------------- 1 | package repositories 2 | 3 | import ( 4 | "context" 5 | "path" 6 | "testing" 7 | 8 | mocket "github.com/Selvatico/go-mocket" 9 | "github.com/flyteorg/datacatalog/pkg/repositories/config" 10 | "github.com/flyteorg/flytestdlib/database" 11 | "github.com/stretchr/testify/assert" 12 | 13 | "database/sql/driver" 14 | 15 | "github.com/flyteorg/datacatalog/pkg/repositories/utils" 16 | ) 17 | 18 | func TestCreateDB(t *testing.T) { 19 | GlobalMock := mocket.Catcher.Reset() 20 | GlobalMock.Logging = true 21 | 22 | checkExists := false 23 | GlobalMock.NewMock().WithQuery( 24 | `SELECT EXISTS(SELECT datname FROM pg_catalog.pg_database WHERE datname = $1)%!(EXTRA string=testDB)`).WithCallback( 25 | func(s string, values []driver.NamedValue) { 26 | checkExists = true 27 | }, 28 | ).WithReply([]map[string]interface{}{ 29 | {"exists": false}, 30 | }) 31 | 32 | createdDB := false 33 | 34 | // NOTE: unfortunately mocket does not support checking CREATE statements, but let's match the suffix 35 | GlobalMock.NewMock().WithQuery( 36 | `DATABASE testDB`).WithCallback( 37 | func(s string, values []driver.NamedValue) { 38 | assert.Equal(t, "CREATE DATABASE testDB", s) 39 | createdDB = true 40 | }, 41 | ) 42 | 43 | db := utils.GetDbForTest(t) 44 | dbHandle := &DBHandle{ 45 | db: db, 46 | } 47 | _ = dbHandle.CreateDB("testDB") 48 | assert.True(t, checkExists) 49 | assert.True(t, createdDB) 50 | } 51 | 52 | func TestDBAlreadyExists(t *testing.T) { 53 | GlobalMock := mocket.Catcher.Reset() 54 | GlobalMock.Logging = true 55 | 56 | checkExists := false 57 | GlobalMock.NewMock().WithQuery( 58 | `SELECT EXISTS(SELECT datname FROM pg_catalog.pg_database WHERE datname = $1)%!(EXTRA string=testDB)`).WithCallback( 59 | func(s string, values []driver.NamedValue) { 60 | checkExists = true 61 | }, 62 | ).WithReply([]map[string]interface{}{ 63 | {"exists": true}, 64 | }) 65 | 66 | createdDB := false 67 | GlobalMock.NewMock().WithQuery( 68 | `DATABASE testDB`).WithCallback( 69 | func(s string, values []driver.NamedValue) { 70 | createdDB = false 71 | }, 72 | ) 73 | 74 | db := utils.GetDbForTest(t) 75 | dbHandle := &DBHandle{ 76 | db: db, 77 | } 78 | err := dbHandle.CreateDB("testDB") 79 | assert.NoError(t, err) 80 | assert.True(t, checkExists) 81 | assert.False(t, createdDB) 82 | } 83 | 84 | func TestNewDBHandle(t *testing.T) { 85 | t.Run("missing DB Config", func(t *testing.T) { 86 | _, err := NewDBHandle(context.TODO(), database.DbConfig{}, migrateScope) 87 | assert.Error(t, err) 88 | }) 89 | 90 | t.Run("sqlite config", func(t *testing.T) { 91 | dbFile := path.Join(t.TempDir(), "admin.db") 92 | dbHandle, err := NewDBHandle(context.TODO(), database.DbConfig{SQLite: database.SQLiteConfig{File: dbFile}}, migrateScope) 93 | assert.Nil(t, err) 94 | assert.NotNil(t, dbHandle) 95 | assert.Equal(t, config.Sqlite, dbHandle.db.Name()) 96 | }) 97 | } 98 | -------------------------------------------------------------------------------- /pkg/repositories/initialize.go: -------------------------------------------------------------------------------- 1 | package repositories 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "reflect" 7 | 8 | errors2 "github.com/flyteorg/datacatalog/pkg/repositories/errors" 9 | "github.com/flyteorg/datacatalog/pkg/runtime" 10 | "github.com/flyteorg/flytestdlib/logger" 11 | "github.com/flyteorg/flytestdlib/promutils" 12 | "github.com/jackc/pgconn" 13 | ) 14 | 15 | var migrationsScope = promutils.NewScope("migrations") 16 | var migrateScope = migrationsScope.NewSubScope("migrate") 17 | 18 | // all postgres servers come by default with a db name named postgres 19 | const defaultDB = "postgres" 20 | const pqInvalidDBCode = "3D000" 21 | const pqDbAlreadyExistsCode = "42P04" 22 | 23 | // Migrate This command will run all the migrations for the database 24 | func Migrate(ctx context.Context) error { 25 | configProvider := runtime.NewConfigurationProvider() 26 | dbConfigValues := *configProvider.ApplicationConfiguration().GetDbConfig() 27 | 28 | dbName := dbConfigValues.Postgres.DbName 29 | dbHandle, err := NewDBHandle(ctx, dbConfigValues, migrateScope) 30 | 31 | if err != nil { 32 | // if db does not exist, try creating it 33 | cErr, ok := err.(errors2.ConnectError) 34 | if !ok { 35 | logger.Errorf(ctx, "Failed to cast error of type: %v, err: %v", reflect.TypeOf(err), 36 | err) 37 | panic(err) 38 | } 39 | pqError := cErr.Unwrap().(*pgconn.PgError) 40 | if pqError.Code == pqInvalidDBCode { 41 | logger.Warningf(ctx, "Database [%v] does not exist, trying to create it now", dbName) 42 | 43 | dbConfigValues.Postgres.DbName = defaultDB 44 | setupDBHandler, err := NewDBHandle(ctx, dbConfigValues, migrateScope) 45 | if err != nil { 46 | logger.Errorf(ctx, "Failed to connect to default DB %v, err %v", defaultDB, err) 47 | panic(err) 48 | } 49 | 50 | // Create the database if it doesn't exist 51 | // NOTE: this is non-destructive - if for some reason one does exist an err will be thrown 52 | err = setupDBHandler.CreateDB(dbName) 53 | if err != nil { 54 | logger.Errorf(ctx, "Failed to create DB %v err %v", dbName, err) 55 | panic(err) 56 | } 57 | 58 | dbConfigValues.Postgres.DbName = dbName 59 | dbHandle, err = NewDBHandle(ctx, dbConfigValues, migrateScope) 60 | if err != nil { 61 | logger.Errorf(ctx, "Failed to connect DB err %v", err) 62 | panic(err) 63 | } 64 | } else { 65 | logger.Errorf(ctx, "Failed to connect DB err %v", err) 66 | panic(err) 67 | } 68 | } 69 | 70 | logger.Infof(ctx, "Created DB connection.") 71 | 72 | // TODO: checkpoints for migrations 73 | if err := dbHandle.Migrate(ctx); err != nil { 74 | logger.Errorf(ctx, "Failed to migrate. err: %v", err) 75 | panic(err) 76 | } 77 | logger.Infof(ctx, "Ran DB migration successfully.") 78 | return nil 79 | } 80 | 81 | func isPgErrorWithCode(err error, code string) bool { 82 | pgErr := &pgconn.PgError{} 83 | if !errors.As(err, &pgErr) { 84 | // err chain does not contain a pgconn.PgError 85 | return false 86 | } 87 | 88 | // pgconn.PgError found in chain and set to code specified 89 | return pgErr.Code == code 90 | } 91 | -------------------------------------------------------------------------------- /pkg/repositories/interfaces/artifact_repo.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | ) 8 | 9 | //go:generate mockery -name=ArtifactRepo -output=../mocks -case=underscore 10 | 11 | type ArtifactRepo interface { 12 | Create(ctx context.Context, in models.Artifact) error 13 | Get(ctx context.Context, in models.ArtifactKey) (models.Artifact, error) 14 | List(ctx context.Context, datasetKey models.DatasetKey, in models.ListModelsInput) ([]models.Artifact, error) 15 | Update(ctx context.Context, artifact models.Artifact) error 16 | } 17 | -------------------------------------------------------------------------------- /pkg/repositories/interfaces/base.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | type DataCatalogRepo interface { 4 | DatasetRepo() DatasetRepo 5 | ArtifactRepo() ArtifactRepo 6 | TagRepo() TagRepo 7 | ReservationRepo() ReservationRepo 8 | } 9 | -------------------------------------------------------------------------------- /pkg/repositories/interfaces/dataset_repo.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | ) 8 | 9 | //go:generate mockery -name=DatasetRepo -output=../mocks -case=underscore 10 | 11 | type DatasetRepo interface { 12 | Create(ctx context.Context, in models.Dataset) error 13 | Get(ctx context.Context, in models.DatasetKey) (models.Dataset, error) 14 | List(ctx context.Context, in models.ListModelsInput) ([]models.Dataset, error) 15 | } 16 | -------------------------------------------------------------------------------- /pkg/repositories/interfaces/partition_repo.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | ) 8 | 9 | //go:generate mockery -name=PartitionRepo -output=../mocks -case=underscore 10 | 11 | type PartitionRepo interface { 12 | Create(ctx context.Context, in models.Partition) error 13 | } 14 | -------------------------------------------------------------------------------- /pkg/repositories/interfaces/reservation_repo.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 8 | ) 9 | 10 | //go:generate mockery -name=ReservationRepo -output=../mocks -case=underscore 11 | 12 | // Interface to interact with Reservation Table 13 | type ReservationRepo interface { 14 | 15 | // Create a new reservation if the reservation does not already exist 16 | Create(ctx context.Context, reservation models.Reservation, now time.Time) error 17 | 18 | // Delete a reservation if it exists 19 | Delete(ctx context.Context, reservation models.ReservationKey, ownerID string) error 20 | 21 | // Get reservation 22 | Get(ctx context.Context, reservationKey models.ReservationKey) (models.Reservation, error) 23 | 24 | // Update an existing reservation. If called by the current owner, we update the 25 | // expiresAt timestamp. If called by a new owner and the current reservation has 26 | // expired, we attempt to take over the reservation. 27 | Update(ctx context.Context, reservation models.Reservation, now time.Time) error 28 | } 29 | -------------------------------------------------------------------------------- /pkg/repositories/interfaces/tag_repo.go: -------------------------------------------------------------------------------- 1 | package interfaces 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | ) 8 | 9 | //go:generate mockery -name=TagRepo -output=../mocks -case=underscore 10 | 11 | type TagRepo interface { 12 | Create(ctx context.Context, in models.Tag) error 13 | Get(ctx context.Context, in models.TagKey) (models.Tag, error) 14 | } 15 | -------------------------------------------------------------------------------- /pkg/repositories/mocks/artifact_repo.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | mock "github.com/stretchr/testify/mock" 9 | 10 | models "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | ) 12 | 13 | // ArtifactRepo is an autogenerated mock type for the ArtifactRepo type 14 | type ArtifactRepo struct { 15 | mock.Mock 16 | } 17 | 18 | type ArtifactRepo_Create struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m ArtifactRepo_Create) Return(_a0 error) *ArtifactRepo_Create { 23 | return &ArtifactRepo_Create{Call: _m.Call.Return(_a0)} 24 | } 25 | 26 | func (_m *ArtifactRepo) OnCreate(ctx context.Context, in models.Artifact) *ArtifactRepo_Create { 27 | c_call := _m.On("Create", ctx, in) 28 | return &ArtifactRepo_Create{Call: c_call} 29 | } 30 | 31 | func (_m *ArtifactRepo) OnCreateMatch(matchers ...interface{}) *ArtifactRepo_Create { 32 | c_call := _m.On("Create", matchers...) 33 | return &ArtifactRepo_Create{Call: c_call} 34 | } 35 | 36 | // Create provides a mock function with given fields: ctx, in 37 | func (_m *ArtifactRepo) Create(ctx context.Context, in models.Artifact) error { 38 | ret := _m.Called(ctx, in) 39 | 40 | var r0 error 41 | if rf, ok := ret.Get(0).(func(context.Context, models.Artifact) error); ok { 42 | r0 = rf(ctx, in) 43 | } else { 44 | r0 = ret.Error(0) 45 | } 46 | 47 | return r0 48 | } 49 | 50 | type ArtifactRepo_Get struct { 51 | *mock.Call 52 | } 53 | 54 | func (_m ArtifactRepo_Get) Return(_a0 models.Artifact, _a1 error) *ArtifactRepo_Get { 55 | return &ArtifactRepo_Get{Call: _m.Call.Return(_a0, _a1)} 56 | } 57 | 58 | func (_m *ArtifactRepo) OnGet(ctx context.Context, in models.ArtifactKey) *ArtifactRepo_Get { 59 | c_call := _m.On("Get", ctx, in) 60 | return &ArtifactRepo_Get{Call: c_call} 61 | } 62 | 63 | func (_m *ArtifactRepo) OnGetMatch(matchers ...interface{}) *ArtifactRepo_Get { 64 | c_call := _m.On("Get", matchers...) 65 | return &ArtifactRepo_Get{Call: c_call} 66 | } 67 | 68 | // Get provides a mock function with given fields: ctx, in 69 | func (_m *ArtifactRepo) Get(ctx context.Context, in models.ArtifactKey) (models.Artifact, error) { 70 | ret := _m.Called(ctx, in) 71 | 72 | var r0 models.Artifact 73 | if rf, ok := ret.Get(0).(func(context.Context, models.ArtifactKey) models.Artifact); ok { 74 | r0 = rf(ctx, in) 75 | } else { 76 | r0 = ret.Get(0).(models.Artifact) 77 | } 78 | 79 | var r1 error 80 | if rf, ok := ret.Get(1).(func(context.Context, models.ArtifactKey) error); ok { 81 | r1 = rf(ctx, in) 82 | } else { 83 | r1 = ret.Error(1) 84 | } 85 | 86 | return r0, r1 87 | } 88 | 89 | type ArtifactRepo_List struct { 90 | *mock.Call 91 | } 92 | 93 | func (_m ArtifactRepo_List) Return(_a0 []models.Artifact, _a1 error) *ArtifactRepo_List { 94 | return &ArtifactRepo_List{Call: _m.Call.Return(_a0, _a1)} 95 | } 96 | 97 | func (_m *ArtifactRepo) OnList(ctx context.Context, datasetKey models.DatasetKey, in models.ListModelsInput) *ArtifactRepo_List { 98 | c_call := _m.On("List", ctx, datasetKey, in) 99 | return &ArtifactRepo_List{Call: c_call} 100 | } 101 | 102 | func (_m *ArtifactRepo) OnListMatch(matchers ...interface{}) *ArtifactRepo_List { 103 | c_call := _m.On("List", matchers...) 104 | return &ArtifactRepo_List{Call: c_call} 105 | } 106 | 107 | // List provides a mock function with given fields: ctx, datasetKey, in 108 | func (_m *ArtifactRepo) List(ctx context.Context, datasetKey models.DatasetKey, in models.ListModelsInput) ([]models.Artifact, error) { 109 | ret := _m.Called(ctx, datasetKey, in) 110 | 111 | var r0 []models.Artifact 112 | if rf, ok := ret.Get(0).(func(context.Context, models.DatasetKey, models.ListModelsInput) []models.Artifact); ok { 113 | r0 = rf(ctx, datasetKey, in) 114 | } else { 115 | if ret.Get(0) != nil { 116 | r0 = ret.Get(0).([]models.Artifact) 117 | } 118 | } 119 | 120 | var r1 error 121 | if rf, ok := ret.Get(1).(func(context.Context, models.DatasetKey, models.ListModelsInput) error); ok { 122 | r1 = rf(ctx, datasetKey, in) 123 | } else { 124 | r1 = ret.Error(1) 125 | } 126 | 127 | return r0, r1 128 | } 129 | 130 | type ArtifactRepo_Update struct { 131 | *mock.Call 132 | } 133 | 134 | func (_m ArtifactRepo_Update) Return(_a0 error) *ArtifactRepo_Update { 135 | return &ArtifactRepo_Update{Call: _m.Call.Return(_a0)} 136 | } 137 | 138 | func (_m *ArtifactRepo) OnUpdate(ctx context.Context, artifact models.Artifact) *ArtifactRepo_Update { 139 | c_call := _m.On("Update", ctx, artifact) 140 | return &ArtifactRepo_Update{Call: c_call} 141 | } 142 | 143 | func (_m *ArtifactRepo) OnUpdateMatch(matchers ...interface{}) *ArtifactRepo_Update { 144 | c_call := _m.On("Update", matchers...) 145 | return &ArtifactRepo_Update{Call: c_call} 146 | } 147 | 148 | // Update provides a mock function with given fields: ctx, artifact 149 | func (_m *ArtifactRepo) Update(ctx context.Context, artifact models.Artifact) error { 150 | ret := _m.Called(ctx, artifact) 151 | 152 | var r0 error 153 | if rf, ok := ret.Get(0).(func(context.Context, models.Artifact) error); ok { 154 | r0 = rf(ctx, artifact) 155 | } else { 156 | r0 = ret.Error(0) 157 | } 158 | 159 | return r0 160 | } 161 | -------------------------------------------------------------------------------- /pkg/repositories/mocks/base.go: -------------------------------------------------------------------------------- 1 | package mocks 2 | 3 | import "github.com/flyteorg/datacatalog/pkg/repositories/interfaces" 4 | 5 | type DataCatalogRepo struct { 6 | MockDatasetRepo *DatasetRepo 7 | MockArtifactRepo *ArtifactRepo 8 | MockTagRepo *TagRepo 9 | MockReservationRepo *ReservationRepo 10 | } 11 | 12 | func (m *DataCatalogRepo) DatasetRepo() interfaces.DatasetRepo { 13 | return m.MockDatasetRepo 14 | } 15 | 16 | func (m *DataCatalogRepo) ArtifactRepo() interfaces.ArtifactRepo { 17 | return m.MockArtifactRepo 18 | } 19 | 20 | func (m *DataCatalogRepo) TagRepo() interfaces.TagRepo { 21 | return m.MockTagRepo 22 | } 23 | 24 | func (m *DataCatalogRepo) ReservationRepo() interfaces.ReservationRepo { 25 | return m.MockReservationRepo 26 | } 27 | -------------------------------------------------------------------------------- /pkg/repositories/mocks/dataset_repo.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | mock "github.com/stretchr/testify/mock" 9 | 10 | models "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | ) 12 | 13 | // DatasetRepo is an autogenerated mock type for the DatasetRepo type 14 | type DatasetRepo struct { 15 | mock.Mock 16 | } 17 | 18 | type DatasetRepo_Create struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m DatasetRepo_Create) Return(_a0 error) *DatasetRepo_Create { 23 | return &DatasetRepo_Create{Call: _m.Call.Return(_a0)} 24 | } 25 | 26 | func (_m *DatasetRepo) OnCreate(ctx context.Context, in models.Dataset) *DatasetRepo_Create { 27 | c_call := _m.On("Create", ctx, in) 28 | return &DatasetRepo_Create{Call: c_call} 29 | } 30 | 31 | func (_m *DatasetRepo) OnCreateMatch(matchers ...interface{}) *DatasetRepo_Create { 32 | c_call := _m.On("Create", matchers...) 33 | return &DatasetRepo_Create{Call: c_call} 34 | } 35 | 36 | // Create provides a mock function with given fields: ctx, in 37 | func (_m *DatasetRepo) Create(ctx context.Context, in models.Dataset) error { 38 | ret := _m.Called(ctx, in) 39 | 40 | var r0 error 41 | if rf, ok := ret.Get(0).(func(context.Context, models.Dataset) error); ok { 42 | r0 = rf(ctx, in) 43 | } else { 44 | r0 = ret.Error(0) 45 | } 46 | 47 | return r0 48 | } 49 | 50 | type DatasetRepo_Get struct { 51 | *mock.Call 52 | } 53 | 54 | func (_m DatasetRepo_Get) Return(_a0 models.Dataset, _a1 error) *DatasetRepo_Get { 55 | return &DatasetRepo_Get{Call: _m.Call.Return(_a0, _a1)} 56 | } 57 | 58 | func (_m *DatasetRepo) OnGet(ctx context.Context, in models.DatasetKey) *DatasetRepo_Get { 59 | c_call := _m.On("Get", ctx, in) 60 | return &DatasetRepo_Get{Call: c_call} 61 | } 62 | 63 | func (_m *DatasetRepo) OnGetMatch(matchers ...interface{}) *DatasetRepo_Get { 64 | c_call := _m.On("Get", matchers...) 65 | return &DatasetRepo_Get{Call: c_call} 66 | } 67 | 68 | // Get provides a mock function with given fields: ctx, in 69 | func (_m *DatasetRepo) Get(ctx context.Context, in models.DatasetKey) (models.Dataset, error) { 70 | ret := _m.Called(ctx, in) 71 | 72 | var r0 models.Dataset 73 | if rf, ok := ret.Get(0).(func(context.Context, models.DatasetKey) models.Dataset); ok { 74 | r0 = rf(ctx, in) 75 | } else { 76 | r0 = ret.Get(0).(models.Dataset) 77 | } 78 | 79 | var r1 error 80 | if rf, ok := ret.Get(1).(func(context.Context, models.DatasetKey) error); ok { 81 | r1 = rf(ctx, in) 82 | } else { 83 | r1 = ret.Error(1) 84 | } 85 | 86 | return r0, r1 87 | } 88 | 89 | type DatasetRepo_List struct { 90 | *mock.Call 91 | } 92 | 93 | func (_m DatasetRepo_List) Return(_a0 []models.Dataset, _a1 error) *DatasetRepo_List { 94 | return &DatasetRepo_List{Call: _m.Call.Return(_a0, _a1)} 95 | } 96 | 97 | func (_m *DatasetRepo) OnList(ctx context.Context, in models.ListModelsInput) *DatasetRepo_List { 98 | c_call := _m.On("List", ctx, in) 99 | return &DatasetRepo_List{Call: c_call} 100 | } 101 | 102 | func (_m *DatasetRepo) OnListMatch(matchers ...interface{}) *DatasetRepo_List { 103 | c_call := _m.On("List", matchers...) 104 | return &DatasetRepo_List{Call: c_call} 105 | } 106 | 107 | // List provides a mock function with given fields: ctx, in 108 | func (_m *DatasetRepo) List(ctx context.Context, in models.ListModelsInput) ([]models.Dataset, error) { 109 | ret := _m.Called(ctx, in) 110 | 111 | var r0 []models.Dataset 112 | if rf, ok := ret.Get(0).(func(context.Context, models.ListModelsInput) []models.Dataset); ok { 113 | r0 = rf(ctx, in) 114 | } else { 115 | if ret.Get(0) != nil { 116 | r0 = ret.Get(0).([]models.Dataset) 117 | } 118 | } 119 | 120 | var r1 error 121 | if rf, ok := ret.Get(1).(func(context.Context, models.ListModelsInput) error); ok { 122 | r1 = rf(ctx, in) 123 | } else { 124 | r1 = ret.Error(1) 125 | } 126 | 127 | return r0, r1 128 | } 129 | -------------------------------------------------------------------------------- /pkg/repositories/mocks/partition_repo.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | mock "github.com/stretchr/testify/mock" 9 | 10 | models "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | ) 12 | 13 | // PartitionRepo is an autogenerated mock type for the PartitionRepo type 14 | type PartitionRepo struct { 15 | mock.Mock 16 | } 17 | 18 | type PartitionRepo_Create struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m PartitionRepo_Create) Return(_a0 error) *PartitionRepo_Create { 23 | return &PartitionRepo_Create{Call: _m.Call.Return(_a0)} 24 | } 25 | 26 | func (_m *PartitionRepo) OnCreate(ctx context.Context, in models.Partition) *PartitionRepo_Create { 27 | c_call := _m.On("Create", ctx, in) 28 | return &PartitionRepo_Create{Call: c_call} 29 | } 30 | 31 | func (_m *PartitionRepo) OnCreateMatch(matchers ...interface{}) *PartitionRepo_Create { 32 | c_call := _m.On("Create", matchers...) 33 | return &PartitionRepo_Create{Call: c_call} 34 | } 35 | 36 | // Create provides a mock function with given fields: ctx, in 37 | func (_m *PartitionRepo) Create(ctx context.Context, in models.Partition) error { 38 | ret := _m.Called(ctx, in) 39 | 40 | var r0 error 41 | if rf, ok := ret.Get(0).(func(context.Context, models.Partition) error); ok { 42 | r0 = rf(ctx, in) 43 | } else { 44 | r0 = ret.Error(0) 45 | } 46 | 47 | return r0 48 | } 49 | -------------------------------------------------------------------------------- /pkg/repositories/mocks/reservation_repo.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | mock "github.com/stretchr/testify/mock" 9 | 10 | models "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | 12 | time "time" 13 | ) 14 | 15 | // ReservationRepo is an autogenerated mock type for the ReservationRepo type 16 | type ReservationRepo struct { 17 | mock.Mock 18 | } 19 | 20 | type ReservationRepo_Create struct { 21 | *mock.Call 22 | } 23 | 24 | func (_m ReservationRepo_Create) Return(_a0 error) *ReservationRepo_Create { 25 | return &ReservationRepo_Create{Call: _m.Call.Return(_a0)} 26 | } 27 | 28 | func (_m *ReservationRepo) OnCreate(ctx context.Context, reservation models.Reservation, now time.Time) *ReservationRepo_Create { 29 | c_call := _m.On("Create", ctx, reservation, now) 30 | return &ReservationRepo_Create{Call: c_call} 31 | } 32 | 33 | func (_m *ReservationRepo) OnCreateMatch(matchers ...interface{}) *ReservationRepo_Create { 34 | c_call := _m.On("Create", matchers...) 35 | return &ReservationRepo_Create{Call: c_call} 36 | } 37 | 38 | // Create provides a mock function with given fields: ctx, reservation, now 39 | func (_m *ReservationRepo) Create(ctx context.Context, reservation models.Reservation, now time.Time) error { 40 | ret := _m.Called(ctx, reservation, now) 41 | 42 | var r0 error 43 | if rf, ok := ret.Get(0).(func(context.Context, models.Reservation, time.Time) error); ok { 44 | r0 = rf(ctx, reservation, now) 45 | } else { 46 | r0 = ret.Error(0) 47 | } 48 | 49 | return r0 50 | } 51 | 52 | type ReservationRepo_Delete struct { 53 | *mock.Call 54 | } 55 | 56 | func (_m ReservationRepo_Delete) Return(_a0 error) *ReservationRepo_Delete { 57 | return &ReservationRepo_Delete{Call: _m.Call.Return(_a0)} 58 | } 59 | 60 | func (_m *ReservationRepo) OnDelete(ctx context.Context, reservation models.ReservationKey, ownerID string) *ReservationRepo_Delete { 61 | c_call := _m.On("Delete", ctx, reservation, ownerID) 62 | return &ReservationRepo_Delete{Call: c_call} 63 | } 64 | 65 | func (_m *ReservationRepo) OnDeleteMatch(matchers ...interface{}) *ReservationRepo_Delete { 66 | c_call := _m.On("Delete", matchers...) 67 | return &ReservationRepo_Delete{Call: c_call} 68 | } 69 | 70 | // Delete provides a mock function with given fields: ctx, reservation, ownerID 71 | func (_m *ReservationRepo) Delete(ctx context.Context, reservation models.ReservationKey, ownerID string) error { 72 | ret := _m.Called(ctx, reservation, ownerID) 73 | 74 | var r0 error 75 | if rf, ok := ret.Get(0).(func(context.Context, models.ReservationKey, string) error); ok { 76 | r0 = rf(ctx, reservation, ownerID) 77 | } else { 78 | r0 = ret.Error(0) 79 | } 80 | 81 | return r0 82 | } 83 | 84 | type ReservationRepo_Get struct { 85 | *mock.Call 86 | } 87 | 88 | func (_m ReservationRepo_Get) Return(_a0 models.Reservation, _a1 error) *ReservationRepo_Get { 89 | return &ReservationRepo_Get{Call: _m.Call.Return(_a0, _a1)} 90 | } 91 | 92 | func (_m *ReservationRepo) OnGet(ctx context.Context, reservationKey models.ReservationKey) *ReservationRepo_Get { 93 | c_call := _m.On("Get", ctx, reservationKey) 94 | return &ReservationRepo_Get{Call: c_call} 95 | } 96 | 97 | func (_m *ReservationRepo) OnGetMatch(matchers ...interface{}) *ReservationRepo_Get { 98 | c_call := _m.On("Get", matchers...) 99 | return &ReservationRepo_Get{Call: c_call} 100 | } 101 | 102 | // Get provides a mock function with given fields: ctx, reservationKey 103 | func (_m *ReservationRepo) Get(ctx context.Context, reservationKey models.ReservationKey) (models.Reservation, error) { 104 | ret := _m.Called(ctx, reservationKey) 105 | 106 | var r0 models.Reservation 107 | if rf, ok := ret.Get(0).(func(context.Context, models.ReservationKey) models.Reservation); ok { 108 | r0 = rf(ctx, reservationKey) 109 | } else { 110 | r0 = ret.Get(0).(models.Reservation) 111 | } 112 | 113 | var r1 error 114 | if rf, ok := ret.Get(1).(func(context.Context, models.ReservationKey) error); ok { 115 | r1 = rf(ctx, reservationKey) 116 | } else { 117 | r1 = ret.Error(1) 118 | } 119 | 120 | return r0, r1 121 | } 122 | 123 | type ReservationRepo_Update struct { 124 | *mock.Call 125 | } 126 | 127 | func (_m ReservationRepo_Update) Return(_a0 error) *ReservationRepo_Update { 128 | return &ReservationRepo_Update{Call: _m.Call.Return(_a0)} 129 | } 130 | 131 | func (_m *ReservationRepo) OnUpdate(ctx context.Context, reservation models.Reservation, now time.Time) *ReservationRepo_Update { 132 | c_call := _m.On("Update", ctx, reservation, now) 133 | return &ReservationRepo_Update{Call: c_call} 134 | } 135 | 136 | func (_m *ReservationRepo) OnUpdateMatch(matchers ...interface{}) *ReservationRepo_Update { 137 | c_call := _m.On("Update", matchers...) 138 | return &ReservationRepo_Update{Call: c_call} 139 | } 140 | 141 | // Update provides a mock function with given fields: ctx, reservation, now 142 | func (_m *ReservationRepo) Update(ctx context.Context, reservation models.Reservation, now time.Time) error { 143 | ret := _m.Called(ctx, reservation, now) 144 | 145 | var r0 error 146 | if rf, ok := ret.Get(0).(func(context.Context, models.Reservation, time.Time) error); ok { 147 | r0 = rf(ctx, reservation, now) 148 | } else { 149 | r0 = ret.Error(0) 150 | } 151 | 152 | return r0 153 | } 154 | -------------------------------------------------------------------------------- /pkg/repositories/mocks/tag_repo.go: -------------------------------------------------------------------------------- 1 | // Code generated by mockery v1.0.1. DO NOT EDIT. 2 | 3 | package mocks 4 | 5 | import ( 6 | context "context" 7 | 8 | mock "github.com/stretchr/testify/mock" 9 | 10 | models "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | ) 12 | 13 | // TagRepo is an autogenerated mock type for the TagRepo type 14 | type TagRepo struct { 15 | mock.Mock 16 | } 17 | 18 | type TagRepo_Create struct { 19 | *mock.Call 20 | } 21 | 22 | func (_m TagRepo_Create) Return(_a0 error) *TagRepo_Create { 23 | return &TagRepo_Create{Call: _m.Call.Return(_a0)} 24 | } 25 | 26 | func (_m *TagRepo) OnCreate(ctx context.Context, in models.Tag) *TagRepo_Create { 27 | c_call := _m.On("Create", ctx, in) 28 | return &TagRepo_Create{Call: c_call} 29 | } 30 | 31 | func (_m *TagRepo) OnCreateMatch(matchers ...interface{}) *TagRepo_Create { 32 | c_call := _m.On("Create", matchers...) 33 | return &TagRepo_Create{Call: c_call} 34 | } 35 | 36 | // Create provides a mock function with given fields: ctx, in 37 | func (_m *TagRepo) Create(ctx context.Context, in models.Tag) error { 38 | ret := _m.Called(ctx, in) 39 | 40 | var r0 error 41 | if rf, ok := ret.Get(0).(func(context.Context, models.Tag) error); ok { 42 | r0 = rf(ctx, in) 43 | } else { 44 | r0 = ret.Error(0) 45 | } 46 | 47 | return r0 48 | } 49 | 50 | type TagRepo_Get struct { 51 | *mock.Call 52 | } 53 | 54 | func (_m TagRepo_Get) Return(_a0 models.Tag, _a1 error) *TagRepo_Get { 55 | return &TagRepo_Get{Call: _m.Call.Return(_a0, _a1)} 56 | } 57 | 58 | func (_m *TagRepo) OnGet(ctx context.Context, in models.TagKey) *TagRepo_Get { 59 | c_call := _m.On("Get", ctx, in) 60 | return &TagRepo_Get{Call: c_call} 61 | } 62 | 63 | func (_m *TagRepo) OnGetMatch(matchers ...interface{}) *TagRepo_Get { 64 | c_call := _m.On("Get", matchers...) 65 | return &TagRepo_Get{Call: c_call} 66 | } 67 | 68 | // Get provides a mock function with given fields: ctx, in 69 | func (_m *TagRepo) Get(ctx context.Context, in models.TagKey) (models.Tag, error) { 70 | ret := _m.Called(ctx, in) 71 | 72 | var r0 models.Tag 73 | if rf, ok := ret.Get(0).(func(context.Context, models.TagKey) models.Tag); ok { 74 | r0 = rf(ctx, in) 75 | } else { 76 | r0 = ret.Get(0).(models.Tag) 77 | } 78 | 79 | var r1 error 80 | if rf, ok := ret.Get(1).(func(context.Context, models.TagKey) error); ok { 81 | r1 = rf(ctx, in) 82 | } else { 83 | r1 = ret.Error(1) 84 | } 85 | 86 | return r0, r1 87 | } 88 | -------------------------------------------------------------------------------- /pkg/repositories/models/artifact.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | type ArtifactKey struct { 4 | DatasetProject string `gorm:"primary_key"` 5 | DatasetName string `gorm:"primary_key"` 6 | DatasetDomain string `gorm:"primary_key"` 7 | DatasetVersion string `gorm:"primary_key"` 8 | ArtifactID string `gorm:"primary_key"` 9 | } 10 | 11 | type Artifact struct { 12 | BaseModel 13 | ArtifactKey 14 | DatasetUUID string `gorm:"type:uuid;index:artifacts_dataset_uuid_idx"` 15 | Dataset Dataset `gorm:"association_autocreate:false"` 16 | ArtifactData []ArtifactData `gorm:"references:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID;foreignkey:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID"` 17 | Partitions []Partition `gorm:"references:ArtifactID;foreignkey:ArtifactID"` 18 | Tags []Tag `gorm:"references:ArtifactID,DatasetUUID;foreignkey:ArtifactID,DatasetUUID"` 19 | SerializedMetadata []byte 20 | } 21 | 22 | type ArtifactData struct { 23 | BaseModel 24 | ArtifactKey 25 | Name string `gorm:"primary_key"` 26 | Location string 27 | } 28 | -------------------------------------------------------------------------------- /pkg/repositories/models/base.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import "time" 4 | 5 | type BaseModel struct { 6 | CreatedAt time.Time 7 | UpdatedAt time.Time 8 | DeletedAt *time.Time `sql:"index"` 9 | } 10 | -------------------------------------------------------------------------------- /pkg/repositories/models/dataset.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import ( 4 | "github.com/gofrs/uuid" 5 | "gorm.io/gorm" 6 | ) 7 | 8 | type DatasetKey struct { 9 | Project string `gorm:"primary_key;"` // part of pkey, no index needed as it is first column in the pkey 10 | Name string `gorm:"primary_key;index:dataset_name_idx"` // part of pkey and has separate index for filtering 11 | Domain string `gorm:"primary_key;index:dataset_domain_idx"` // part of pkey and has separate index for filtering 12 | Version string `gorm:"primary_key;index:dataset_version_idx"` // part of pkey and has separate index for filtering 13 | UUID string `gorm:"type:uuid;unique;"` 14 | } 15 | 16 | type Dataset struct { 17 | BaseModel 18 | DatasetKey 19 | SerializedMetadata []byte 20 | PartitionKeys []PartitionKey `gorm:"references:UUID;foreignkey:DatasetUUID"` 21 | } 22 | 23 | type PartitionKey struct { 24 | BaseModel 25 | DatasetUUID string `gorm:"type:uuid;primary_key"` 26 | Name string `gorm:"primary_key"` 27 | } 28 | 29 | // BeforeCreate so that we set the UUID in golang rather than from a DB function call 30 | func (dataset *Dataset) BeforeCreate(tx *gorm.DB) error { 31 | if dataset.UUID == "" { 32 | generated, err := uuid.NewV4() 33 | if err != nil { 34 | return err 35 | } 36 | 37 | tx.Model(dataset).Update("UUID", generated) 38 | } 39 | return nil 40 | } 41 | -------------------------------------------------------------------------------- /pkg/repositories/models/list.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import "github.com/flyteorg/datacatalog/pkg/common" 4 | 5 | // Inputs to specify to list models 6 | type ListModelsInput struct { 7 | // The filters for the list 8 | ModelFilters []ModelFilter 9 | // The number of models to list 10 | Limit int 11 | // The token to offset results by 12 | Offset int 13 | // Parameter to sort by 14 | SortParameter SortParameter 15 | } 16 | 17 | type SortParameter interface { 18 | GetDBOrderExpression(tableName string) string 19 | } 20 | 21 | // Generates db filter expressions for model values 22 | type ModelValueFilter interface { 23 | GetDBQueryExpression(tableName string) (DBQueryExpr, error) 24 | } 25 | 26 | // Generates the join expressions for filters that require other entities 27 | type ModelJoinCondition interface { 28 | GetJoinOnDBQueryExpression(sourceTableName string, joiningTableName string, joiningTableAlias string) (string, error) 29 | } 30 | 31 | // A single filter for a model encompasses value filters and optionally a join condition if the filter is not on 32 | // the source model 33 | type ModelFilter struct { 34 | ValueFilters []ModelValueFilter 35 | JoinCondition ModelJoinCondition 36 | Entity common.Entity 37 | } 38 | 39 | // Encapsulates the query and necessary arguments to issue a DB query. 40 | type DBQueryExpr struct { 41 | Query string 42 | Args interface{} 43 | } 44 | -------------------------------------------------------------------------------- /pkg/repositories/models/partition.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | // Main Use cases: 4 | // 1. Filter artifacts by partition key/val in a dataset from UI [x] 5 | // 2. Get the artifact that has the partitions (x,y,z + tag_name) = latest [x] 6 | type Partition struct { 7 | BaseModel 8 | DatasetUUID string `gorm:"primary_key;type:uuid"` 9 | Key string `gorm:"primary_key"` 10 | Value string `gorm:"primary_key"` 11 | ArtifactID string `gorm:"primary_key;index"` // index for JOINs with the Tag/Labels table when querying artifacts 12 | } 13 | -------------------------------------------------------------------------------- /pkg/repositories/models/reservation.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import "time" 4 | 5 | // ReservationKey uniquely identifies a reservation 6 | type ReservationKey struct { 7 | DatasetProject string `gorm:"primary_key"` 8 | DatasetName string `gorm:"primary_key"` 9 | DatasetDomain string `gorm:"primary_key"` 10 | DatasetVersion string `gorm:"primary_key"` 11 | TagName string `gorm:"primary_key"` 12 | } 13 | 14 | // Reservation tracks the metadata needed to allow 15 | // task cache serialization 16 | type Reservation struct { 17 | BaseModel 18 | ReservationKey 19 | 20 | // Identifies who owns the reservation 21 | OwnerID string 22 | 23 | // When the reservation will expire 24 | ExpiresAt time.Time 25 | SerializedMetadata []byte 26 | } 27 | -------------------------------------------------------------------------------- /pkg/repositories/models/tag.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | type TagKey struct { 4 | DatasetProject string `gorm:"primary_key"` 5 | DatasetName string `gorm:"primary_key"` 6 | DatasetDomain string `gorm:"primary_key"` 7 | DatasetVersion string `gorm:"primary_key"` 8 | TagName string `gorm:"primary_key"` 9 | } 10 | 11 | type Tag struct { 12 | BaseModel 13 | TagKey 14 | ArtifactID string 15 | DatasetUUID string `gorm:"type:uuid;index:tags_dataset_uuid_idx"` 16 | Artifact Artifact `gorm:"references:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID;foreignkey:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID"` 17 | } 18 | -------------------------------------------------------------------------------- /pkg/repositories/postgres_repo.go: -------------------------------------------------------------------------------- 1 | package repositories 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/repositories/errors" 5 | "github.com/flyteorg/datacatalog/pkg/repositories/gormimpl" 6 | "github.com/flyteorg/datacatalog/pkg/repositories/interfaces" 7 | "github.com/flyteorg/flytestdlib/promutils" 8 | "gorm.io/gorm" 9 | ) 10 | 11 | type PostgresRepo struct { 12 | datasetRepo interfaces.DatasetRepo 13 | artifactRepo interfaces.ArtifactRepo 14 | tagRepo interfaces.TagRepo 15 | reservationRepo interfaces.ReservationRepo 16 | } 17 | 18 | func (dc *PostgresRepo) DatasetRepo() interfaces.DatasetRepo { 19 | return dc.datasetRepo 20 | } 21 | 22 | func (dc *PostgresRepo) ArtifactRepo() interfaces.ArtifactRepo { 23 | return dc.artifactRepo 24 | } 25 | 26 | func (dc *PostgresRepo) TagRepo() interfaces.TagRepo { 27 | return dc.tagRepo 28 | } 29 | 30 | func (dc *PostgresRepo) ReservationRepo() interfaces.ReservationRepo { 31 | return dc.reservationRepo 32 | } 33 | 34 | func NewPostgresRepo(db *gorm.DB, errorTransformer errors.ErrorTransformer, scope promutils.Scope) interfaces.DataCatalogRepo { 35 | return &PostgresRepo{ 36 | datasetRepo: gormimpl.NewDatasetRepo(db, errorTransformer, scope.NewSubScope("dataset")), 37 | artifactRepo: gormimpl.NewArtifactRepo(db, errorTransformer, scope.NewSubScope("artifact")), 38 | tagRepo: gormimpl.NewTagRepo(db, errorTransformer, scope.NewSubScope("tag")), 39 | reservationRepo: gormimpl.NewReservationRepo(db, errorTransformer, scope.NewSubScope("reservation")), 40 | } 41 | } 42 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/artifact.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/errors" 5 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 6 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 7 | "github.com/golang/protobuf/ptypes" 8 | "google.golang.org/grpc/codes" 9 | ) 10 | 11 | func CreateArtifactModel(request *datacatalog.CreateArtifactRequest, artifactData []models.ArtifactData, dataset models.Dataset) (models.Artifact, error) { 12 | datasetID := request.Artifact.Dataset 13 | 14 | serializedMetadata, err := marshalMetadata(request.Artifact.Metadata) 15 | if err != nil { 16 | return models.Artifact{}, err 17 | } 18 | 19 | partitions := make([]models.Partition, len(request.Artifact.Partitions)) 20 | for i, partition := range request.Artifact.GetPartitions() { 21 | partitions[i] = models.Partition{ 22 | DatasetUUID: dataset.UUID, 23 | Key: partition.Key, 24 | Value: partition.Value, 25 | } 26 | } 27 | 28 | return models.Artifact{ 29 | ArtifactKey: models.ArtifactKey{ 30 | DatasetProject: datasetID.Project, 31 | DatasetDomain: datasetID.Domain, 32 | DatasetName: datasetID.Name, 33 | DatasetVersion: datasetID.Version, 34 | ArtifactID: request.Artifact.Id, 35 | }, 36 | DatasetUUID: dataset.UUID, 37 | ArtifactData: artifactData, 38 | SerializedMetadata: serializedMetadata, 39 | Partitions: partitions, 40 | }, nil 41 | } 42 | 43 | func FromArtifactModel(artifact models.Artifact) (*datacatalog.Artifact, error) { 44 | datasetID := &datacatalog.DatasetID{ 45 | Project: artifact.DatasetProject, 46 | Domain: artifact.DatasetDomain, 47 | Name: artifact.DatasetName, 48 | Version: artifact.DatasetVersion, 49 | UUID: artifact.DatasetUUID, 50 | } 51 | 52 | metadata, err := unmarshalMetadata(artifact.SerializedMetadata) 53 | if err != nil { 54 | return &datacatalog.Artifact{}, err 55 | } 56 | 57 | partitions := make([]*datacatalog.Partition, len(artifact.Partitions)) 58 | for i, partition := range artifact.Partitions { 59 | partitions[i] = &datacatalog.Partition{ 60 | Key: partition.Key, 61 | Value: partition.Value, 62 | } 63 | } 64 | 65 | tags := make([]*datacatalog.Tag, len(artifact.Tags)) 66 | for i, tag := range artifact.Tags { 67 | tags[i] = FromTagModel(datasetID, tag) 68 | } 69 | 70 | createdAt, err := ptypes.TimestampProto(artifact.CreatedAt) 71 | if err != nil { 72 | return &datacatalog.Artifact{}, errors.NewDataCatalogErrorf(codes.Internal, 73 | "artifact [%+v] invalid createdAt time conversion", artifact) 74 | } 75 | return &datacatalog.Artifact{ 76 | Id: artifact.ArtifactID, 77 | Dataset: datasetID, 78 | Metadata: metadata, 79 | Partitions: partitions, 80 | Tags: tags, 81 | CreatedAt: createdAt, 82 | }, nil 83 | } 84 | 85 | func FromArtifactModels(artifacts []models.Artifact) ([]*datacatalog.Artifact, error) { 86 | retArtifacts := make([]*datacatalog.Artifact, 0, len(artifacts)) 87 | for _, artifact := range artifacts { 88 | retArtifact, err := FromArtifactModel(artifact) 89 | if err != nil { 90 | return nil, err 91 | } 92 | 93 | retArtifacts = append(retArtifacts, retArtifact) 94 | } 95 | 96 | return retArtifacts, nil 97 | } 98 | 99 | // Transforms datasetID and artifact combination into an ArtifactKey 100 | // The DatasetID is optional since artifactIDs are unique per Artifact 101 | func ToArtifactKey(datasetID *datacatalog.DatasetID, artifactID string) models.ArtifactKey { 102 | artifactKey := models.ArtifactKey{ 103 | ArtifactID: artifactID, 104 | } 105 | if datasetID != nil { 106 | artifactKey.DatasetProject = datasetID.Project 107 | artifactKey.DatasetDomain = datasetID.Domain 108 | artifactKey.DatasetName = datasetID.Name 109 | artifactKey.DatasetVersion = datasetID.Version 110 | } 111 | return artifactKey 112 | } 113 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/dataset.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 5 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 6 | ) 7 | 8 | // Create a dataset model from the Dataset api object. This will serialize the metadata in the dataset as part of the transform 9 | func CreateDatasetModel(dataset *datacatalog.Dataset) (*models.Dataset, error) { 10 | serializedMetadata, err := marshalMetadata(dataset.Metadata) 11 | if err != nil { 12 | return nil, err 13 | } 14 | 15 | partitionKeys := make([]models.PartitionKey, len(dataset.PartitionKeys)) 16 | 17 | for i, partitionKey := range dataset.GetPartitionKeys() { 18 | partitionKeys[i] = models.PartitionKey{ 19 | Name: partitionKey, 20 | } 21 | } 22 | 23 | return &models.Dataset{ 24 | DatasetKey: models.DatasetKey{ 25 | Project: dataset.Id.Project, 26 | Domain: dataset.Id.Domain, 27 | Name: dataset.Id.Name, 28 | Version: dataset.Id.Version, 29 | UUID: dataset.Id.UUID, 30 | }, 31 | SerializedMetadata: serializedMetadata, 32 | PartitionKeys: partitionKeys, 33 | }, nil 34 | } 35 | 36 | // Create a dataset ID from the dataset key model 37 | func FromDatasetID(datasetID *datacatalog.DatasetID) models.DatasetKey { 38 | return models.DatasetKey{ 39 | Project: datasetID.Project, 40 | Domain: datasetID.Domain, 41 | Name: datasetID.Name, 42 | Version: datasetID.Version, 43 | UUID: datasetID.UUID, 44 | } 45 | } 46 | 47 | // Create a Dataset api object given a model, this will unmarshal the metadata into the object as part of the transform 48 | func FromDatasetModel(dataset models.Dataset) (*datacatalog.Dataset, error) { 49 | metadata, err := unmarshalMetadata(dataset.SerializedMetadata) 50 | if err != nil { 51 | return nil, err 52 | } 53 | 54 | partitionKeyStrings := FromPartitionKeyModel(dataset.PartitionKeys) 55 | return &datacatalog.Dataset{ 56 | Id: &datacatalog.DatasetID{ 57 | UUID: dataset.UUID, 58 | Project: dataset.Project, 59 | Domain: dataset.Domain, 60 | Name: dataset.Name, 61 | Version: dataset.Version, 62 | }, 63 | Metadata: metadata, 64 | PartitionKeys: partitionKeyStrings, 65 | }, nil 66 | } 67 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/dataset_test.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | var metadata = &datacatalog.Metadata{ 12 | KeyMap: map[string]string{ 13 | "testKey1": "testValue1", 14 | "testKey2": "testValue2", 15 | }, 16 | } 17 | 18 | var datasetID = &datacatalog.DatasetID{ 19 | Project: "test-project", 20 | Domain: "test-domain", 21 | Name: "test-name", 22 | Version: "test-version", 23 | UUID: "test-uuid", 24 | } 25 | 26 | func assertDatasetIDEqualsModel(t *testing.T, idlDataset *datacatalog.DatasetID, model *models.DatasetKey) { 27 | assert.Equal(t, idlDataset.Project, model.Project) 28 | assert.Equal(t, idlDataset.Domain, model.Domain) 29 | assert.Equal(t, idlDataset.Name, model.Name) 30 | assert.Equal(t, idlDataset.Version, model.Version) 31 | assert.Equal(t, idlDataset.UUID, model.UUID) 32 | } 33 | 34 | func TestCreateDatasetModelNoParitions(t *testing.T) { 35 | dataset := &datacatalog.Dataset{ 36 | Id: datasetID, 37 | Metadata: metadata, 38 | } 39 | 40 | datasetModel, err := CreateDatasetModel(dataset) 41 | assert.NoError(t, err) 42 | assertDatasetIDEqualsModel(t, dataset.Id, &datasetModel.DatasetKey) 43 | 44 | unmarshaledMetadata, err := unmarshalMetadata(datasetModel.SerializedMetadata) 45 | assert.NoError(t, err) 46 | assert.EqualValues(t, unmarshaledMetadata.KeyMap, metadata.KeyMap) 47 | 48 | assert.Len(t, datasetModel.PartitionKeys, 0) 49 | } 50 | 51 | func TestCreateDatasetModel(t *testing.T) { 52 | dataset := &datacatalog.Dataset{ 53 | Id: datasetID, 54 | Metadata: metadata, 55 | PartitionKeys: []string{"key1", "key2"}, 56 | } 57 | 58 | datasetModel, err := CreateDatasetModel(dataset) 59 | assert.NoError(t, err) 60 | assertDatasetIDEqualsModel(t, dataset.Id, &datasetModel.DatasetKey) 61 | 62 | unmarshaledMetadata, err := unmarshalMetadata(datasetModel.SerializedMetadata) 63 | assert.NoError(t, err) 64 | assert.EqualValues(t, unmarshaledMetadata.KeyMap, metadata.KeyMap) 65 | 66 | assert.Len(t, datasetModel.PartitionKeys, 2) 67 | assert.Equal(t, datasetModel.PartitionKeys[0], models.PartitionKey{Name: dataset.PartitionKeys[0]}) 68 | assert.Equal(t, datasetModel.PartitionKeys[1], models.PartitionKey{Name: dataset.PartitionKeys[1]}) 69 | } 70 | 71 | func TestFromDatasetID(t *testing.T) { 72 | datasetKey := FromDatasetID(datasetID) 73 | assertDatasetIDEqualsModel(t, datasetID, &datasetKey) 74 | } 75 | 76 | func TestFromDatasetModelNoPartitionsOrMetadata(t *testing.T) { 77 | datasetModel := &models.Dataset{ 78 | DatasetKey: models.DatasetKey{ 79 | Project: "test-project", 80 | Domain: "test-domain", 81 | Name: "test-name", 82 | Version: "test-version", 83 | }, 84 | SerializedMetadata: []byte{}, 85 | } 86 | dataset, err := FromDatasetModel(*datasetModel) 87 | assert.NoError(t, err) 88 | assertDatasetIDEqualsModel(t, dataset.Id, &datasetModel.DatasetKey) 89 | assert.Len(t, dataset.Metadata.KeyMap, 0) 90 | assert.Len(t, dataset.PartitionKeys, 0) 91 | } 92 | 93 | func TestFromDatasetModelWithPartitions(t *testing.T) { 94 | datasetModel := &models.Dataset{ 95 | DatasetKey: models.DatasetKey{ 96 | Project: "test-project", 97 | Domain: "test-domain", 98 | Name: "test-name", 99 | Version: "test-version", 100 | UUID: "test-uuid", 101 | }, 102 | SerializedMetadata: []byte{10, 22, 10, 8, 116, 101, 115, 116, 75, 101, 121, 49, 18, 10, 116, 101, 115, 116, 86, 97, 108, 117, 101, 49, 10, 22, 10, 8, 116, 101, 115, 116, 75, 101, 121, 50, 18, 10, 116, 101, 115, 116, 86, 97, 108, 117, 101, 50}, 103 | PartitionKeys: []models.PartitionKey{ 104 | {DatasetUUID: "test-uuid", Name: "key1"}, 105 | {DatasetUUID: "test-uuid", Name: "key2"}, 106 | }, 107 | } 108 | dataset, err := FromDatasetModel(*datasetModel) 109 | assert.NoError(t, err) 110 | assertDatasetIDEqualsModel(t, dataset.Id, &datasetModel.DatasetKey) 111 | assert.Len(t, dataset.Metadata.KeyMap, 2) 112 | assert.EqualValues(t, dataset.Metadata.KeyMap, metadata.KeyMap) 113 | assert.Len(t, dataset.PartitionKeys, 2) 114 | } 115 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/filters_test.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/common" 8 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 9 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 10 | "github.com/stretchr/testify/assert" 11 | ) 12 | 13 | func assertJoinExpression(t *testing.T, joinCondition models.ModelJoinCondition, sourceTableName string, joiningTableName string, joiningTableAlias string, expectedJoinStatement string) { 14 | expr, err := joinCondition.GetJoinOnDBQueryExpression(sourceTableName, joiningTableName, joiningTableAlias) 15 | assert.NoError(t, err) 16 | assert.Equal(t, expr, expectedJoinStatement) 17 | } 18 | 19 | func assertFilterExpression(t *testing.T, filter models.ModelValueFilter, tableName string, expectedStatement string, expectedArgs interface{}) { 20 | expr, err := filter.GetDBQueryExpression(tableName) 21 | assert.NoError(t, err) 22 | assert.Equal(t, expectedStatement, expr.Query) 23 | assert.EqualValues(t, expectedArgs, expr.Args) 24 | } 25 | 26 | func TestListInputWithPartitionsAndTags(t *testing.T) { 27 | filter := &datacatalog.FilterExpression{ 28 | Filters: []*datacatalog.SinglePropertyFilter{ 29 | { 30 | PropertyFilter: &datacatalog.SinglePropertyFilter_PartitionFilter{ 31 | PartitionFilter: &datacatalog.PartitionPropertyFilter{ 32 | Property: &datacatalog.PartitionPropertyFilter_KeyVal{ 33 | KeyVal: &datacatalog.KeyValuePair{Key: "key1", Value: "val1"}, 34 | }, 35 | }, 36 | }, 37 | }, 38 | { 39 | PropertyFilter: &datacatalog.SinglePropertyFilter_PartitionFilter{ 40 | PartitionFilter: &datacatalog.PartitionPropertyFilter{ 41 | Property: &datacatalog.PartitionPropertyFilter_KeyVal{ 42 | KeyVal: &datacatalog.KeyValuePair{Key: "key2", Value: "val2"}, 43 | }, 44 | }, 45 | }, 46 | }, 47 | { 48 | PropertyFilter: &datacatalog.SinglePropertyFilter_TagFilter{ 49 | TagFilter: &datacatalog.TagPropertyFilter{ 50 | Property: &datacatalog.TagPropertyFilter_TagName{ 51 | TagName: "special", 52 | }, 53 | }, 54 | }, 55 | }, 56 | }, 57 | } 58 | listInput, err := FilterToListInput(context.Background(), common.Artifact, filter) 59 | assert.NoError(t, err) 60 | 61 | // Should have 3 filters: 2 for partitions, 1 for tag 62 | assert.Len(t, listInput.ModelFilters, 3) 63 | 64 | assertFilterExpression(t, listInput.ModelFilters[0].ValueFilters[0], "partitions", 65 | "partitions.key = ?", "key1") 66 | assertFilterExpression(t, listInput.ModelFilters[0].ValueFilters[1], "partitions", 67 | "partitions.value = ?", "val1") 68 | assertJoinExpression(t, listInput.ModelFilters[0].JoinCondition, "artifacts", "partitions", 69 | "p1", "JOIN partitions p1 ON artifacts.artifact_id = p1.artifact_id") 70 | 71 | assertFilterExpression(t, listInput.ModelFilters[1].ValueFilters[0], "partitions", 72 | "partitions.key = ?", "key2") 73 | assertFilterExpression(t, listInput.ModelFilters[1].ValueFilters[1], "partitions", 74 | "partitions.value = ?", "val2") 75 | assertJoinExpression(t, listInput.ModelFilters[1].JoinCondition, "artifacts", "partitions", 76 | "p2", "JOIN partitions p2 ON artifacts.artifact_id = p2.artifact_id") 77 | 78 | assertFilterExpression(t, listInput.ModelFilters[2].ValueFilters[0], "tags", 79 | "tags.tag_name = ?", "special") 80 | assertJoinExpression(t, listInput.ModelFilters[2].JoinCondition, "artifacts", "tags", 81 | "t1", "JOIN tags t1 ON artifacts.artifact_id = t1.artifact_id") 82 | 83 | } 84 | 85 | func TestEmptyFiledListInput(t *testing.T) { 86 | filter := &datacatalog.FilterExpression{ 87 | Filters: []*datacatalog.SinglePropertyFilter{ 88 | { 89 | PropertyFilter: &datacatalog.SinglePropertyFilter_PartitionFilter{ 90 | PartitionFilter: &datacatalog.PartitionPropertyFilter{ 91 | Property: &datacatalog.PartitionPropertyFilter_KeyVal{ 92 | KeyVal: &datacatalog.KeyValuePair{Key: "", Value: ""}, 93 | }, 94 | }, 95 | }, 96 | }, 97 | }, 98 | } 99 | _, err := FilterToListInput(context.Background(), common.Artifact, filter) 100 | assert.Error(t, err) 101 | } 102 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/pagination.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "strconv" 5 | "strings" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/common" 8 | "github.com/flyteorg/datacatalog/pkg/errors" 9 | "github.com/flyteorg/datacatalog/pkg/repositories/gormimpl" 10 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 11 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 12 | "google.golang.org/grpc/codes" 13 | ) 14 | 15 | func ApplyPagination(paginationOpts *datacatalog.PaginationOptions, input *models.ListModelsInput) error { 16 | var ( 17 | offset = common.DefaultPageOffset 18 | limit = common.MaxPageLimit 19 | sortKey = datacatalog.PaginationOptions_CREATION_TIME 20 | sortOrder = datacatalog.PaginationOptions_DESCENDING 21 | ) 22 | 23 | if paginationOpts != nil { 24 | // if the token is empty, that is still valid input since it is optional 25 | if len(strings.Trim(paginationOpts.Token, " ")) == 0 { 26 | offset = common.DefaultPageOffset 27 | } else { 28 | parsedOffset, err := strconv.ParseInt(paginationOpts.Token, 10, 32) 29 | if err != nil { 30 | return errors.NewDataCatalogErrorf(codes.InvalidArgument, "Invalid token %v", offset) 31 | } 32 | offset = int(parsedOffset) 33 | } 34 | limit = int(paginationOpts.Limit) 35 | sortKey = paginationOpts.SortKey 36 | sortOrder = paginationOpts.SortOrder 37 | } 38 | 39 | input.Offset = offset 40 | input.Limit = limit 41 | input.SortParameter = gormimpl.NewGormSortParameter(sortKey, sortOrder) 42 | return nil 43 | } 44 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/pagination_test.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/common" 7 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 8 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func TestPaginationDefaults(t *testing.T) { 13 | listModelsInput := &models.ListModelsInput{} 14 | err := ApplyPagination(nil, listModelsInput) 15 | assert.NoError(t, err) 16 | assert.Equal(t, common.DefaultPageOffset, listModelsInput.Offset) 17 | assert.Equal(t, common.MaxPageLimit, listModelsInput.Limit) 18 | assert.Equal(t, "artifacts.created_at desc", listModelsInput.SortParameter.GetDBOrderExpression("artifacts")) 19 | } 20 | 21 | func TestPaginationInvalidToken(t *testing.T) { 22 | listModelsInput := &models.ListModelsInput{} 23 | err := ApplyPagination(&datacatalog.PaginationOptions{Token: "pg. 1"}, listModelsInput) 24 | assert.Error(t, err) 25 | } 26 | 27 | func TestCorrectPagination(t *testing.T) { 28 | listModelsInput := &models.ListModelsInput{} 29 | err := ApplyPagination(&datacatalog.PaginationOptions{ 30 | Token: "100", 31 | Limit: 50, 32 | SortKey: datacatalog.PaginationOptions_CREATION_TIME, 33 | SortOrder: datacatalog.PaginationOptions_DESCENDING, 34 | }, listModelsInput) 35 | assert.NoError(t, err) 36 | assert.Equal(t, 50, listModelsInput.Limit) 37 | assert.Equal(t, 100, listModelsInput.Offset) 38 | assert.Equal(t, "artifacts.created_at desc", listModelsInput.SortParameter.GetDBOrderExpression("artifacts")) 39 | } 40 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/partition.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import "github.com/flyteorg/datacatalog/pkg/repositories/models" 4 | 5 | func FromPartitionKeyModel(partitionKeys []models.PartitionKey) []string { 6 | partitionKeyStrings := make([]string, len(partitionKeys)) 7 | for i, partitionKey := range partitionKeys { 8 | partitionKeyStrings[i] = partitionKey.Name 9 | } 10 | 11 | return partitionKeyStrings 12 | } 13 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/reservation.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "time" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/errors" 7 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 8 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | 10 | "github.com/golang/protobuf/ptypes" 11 | 12 | "google.golang.org/grpc/codes" 13 | ) 14 | 15 | func FromReservationID(reservationID *datacatalog.ReservationID) models.ReservationKey { 16 | datasetID := reservationID.DatasetId 17 | 18 | return models.ReservationKey{ 19 | DatasetProject: datasetID.Project, 20 | DatasetDomain: datasetID.Domain, 21 | DatasetName: datasetID.Name, 22 | DatasetVersion: datasetID.Version, 23 | TagName: reservationID.TagName, 24 | } 25 | } 26 | 27 | func CreateReservation(reservation *models.Reservation, heartbeatInterval time.Duration) (datacatalog.Reservation, error) { 28 | expiresAtPb, err := ptypes.TimestampProto(reservation.ExpiresAt) 29 | if err != nil { 30 | return datacatalog.Reservation{}, errors.NewDataCatalogErrorf(codes.Internal, "failed to serialize expires at time") 31 | } 32 | 33 | heartbeatIntervalPb := ptypes.DurationProto(heartbeatInterval) 34 | return datacatalog.Reservation{ 35 | ReservationId: &datacatalog.ReservationID{ 36 | DatasetId: &datacatalog.DatasetID{ 37 | Project: reservation.DatasetProject, 38 | Domain: reservation.DatasetDomain, 39 | Name: reservation.DatasetName, 40 | Version: reservation.DatasetVersion, 41 | }, 42 | TagName: reservation.TagName, 43 | }, 44 | OwnerId: reservation.OwnerID, 45 | HeartbeatInterval: heartbeatIntervalPb, 46 | ExpiresAt: expiresAtPb, 47 | }, nil 48 | } 49 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/reservation_test.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 8 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func TestFromReservationID(t *testing.T) { 13 | reservationID := datacatalog.ReservationID{ 14 | DatasetId: &datacatalog.DatasetID{ 15 | Project: "p", 16 | Name: "n", 17 | Domain: "d", 18 | Version: "v", 19 | }, 20 | TagName: "t", 21 | } 22 | 23 | reservationKey := FromReservationID(&reservationID) 24 | assert.Equal(t, reservationKey.DatasetProject, reservationID.DatasetId.Project) 25 | assert.Equal(t, reservationKey.DatasetName, reservationID.DatasetId.Name) 26 | assert.Equal(t, reservationKey.DatasetDomain, reservationID.DatasetId.Domain) 27 | assert.Equal(t, reservationKey.DatasetVersion, reservationID.DatasetId.Version) 28 | assert.Equal(t, reservationKey.TagName, reservationID.TagName) 29 | } 30 | 31 | func TestCreateReservation(t *testing.T) { 32 | now := time.Now() 33 | heartbeatInterval := time.Second * 5 34 | modelReservation := models.Reservation{ 35 | ReservationKey: models.ReservationKey{ 36 | DatasetProject: "p", 37 | DatasetName: "n", 38 | DatasetDomain: "d", 39 | DatasetVersion: "v", 40 | TagName: "t", 41 | }, 42 | OwnerID: "o", 43 | ExpiresAt: now, 44 | } 45 | 46 | reservation, err := CreateReservation(&modelReservation, heartbeatInterval) 47 | 48 | assert.Equal(t, err, nil) 49 | assert.Equal(t, reservation.ExpiresAt.AsTime(), modelReservation.ExpiresAt.UTC()) 50 | assert.Equal(t, reservation.HeartbeatInterval.AsDuration(), heartbeatInterval) 51 | assert.Equal(t, reservation.OwnerId, modelReservation.OwnerID) 52 | 53 | reservationID := reservation.ReservationId 54 | assert.Equal(t, reservationID.TagName, modelReservation.TagName) 55 | 56 | datasetID := reservationID.DatasetId 57 | assert.Equal(t, datasetID.Project, modelReservation.DatasetProject) 58 | assert.Equal(t, datasetID.Name, modelReservation.DatasetName) 59 | assert.Equal(t, datasetID.Domain, modelReservation.DatasetDomain) 60 | assert.Equal(t, datasetID.Version, modelReservation.DatasetVersion) 61 | } 62 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/tag.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 5 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 6 | ) 7 | 8 | func ToTagKey(datasetID *datacatalog.DatasetID, tagName string) models.TagKey { 9 | return models.TagKey{ 10 | DatasetProject: datasetID.Project, 11 | DatasetDomain: datasetID.Domain, 12 | DatasetName: datasetID.Name, 13 | DatasetVersion: datasetID.Version, 14 | TagName: tagName, 15 | } 16 | } 17 | 18 | func FromTagModel(datasetID *datacatalog.DatasetID, tag models.Tag) *datacatalog.Tag { 19 | return &datacatalog.Tag{ 20 | Name: tag.TagName, 21 | ArtifactId: tag.ArtifactID, 22 | Dataset: datasetID, 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/tag_test.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/flyteorg/datacatalog/pkg/repositories/models" 7 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestToTagKey(t *testing.T) { 12 | datasetID := &datacatalog.DatasetID{ 13 | Project: "testProj", 14 | Domain: "testDomain", 15 | Name: "testName", 16 | Version: "testVersion", 17 | UUID: "test-uuid", 18 | } 19 | 20 | tagName := "testTag" 21 | tagKey := ToTagKey(datasetID, tagName) 22 | 23 | assert.Equal(t, tagName, tagKey.TagName) 24 | assert.Equal(t, datasetID.Project, tagKey.DatasetProject) 25 | assert.Equal(t, datasetID.Domain, tagKey.DatasetDomain) 26 | assert.Equal(t, datasetID.Name, tagKey.DatasetName) 27 | assert.Equal(t, datasetID.Version, tagKey.DatasetVersion) 28 | } 29 | 30 | func TestFromTagModel(t *testing.T) { 31 | datasetID := &datacatalog.DatasetID{ 32 | Project: "testProj", 33 | Domain: "testDomain", 34 | Name: "testName", 35 | Version: "testVersion", 36 | UUID: "test-uuid", 37 | } 38 | 39 | tagModel := models.Tag{ 40 | TagKey: models.TagKey{ 41 | TagName: "test-tag", 42 | }, 43 | DatasetUUID: "dataset-uuid", 44 | } 45 | 46 | tag := FromTagModel(datasetID, tagModel) 47 | 48 | assert.Equal(t, tag.Name, tagModel.TagName) 49 | assert.Equal(t, datasetID.Project, tag.Dataset.Project) 50 | assert.Equal(t, datasetID.Domain, tag.Dataset.Domain) 51 | assert.Equal(t, datasetID.Name, tag.Dataset.Name) 52 | assert.Equal(t, datasetID.Version, tag.Dataset.Version) 53 | assert.Equal(t, datasetID.UUID, tag.Dataset.UUID) 54 | } 55 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/util.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "github.com/flyteorg/datacatalog/pkg/errors" 5 | "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" 6 | "github.com/golang/protobuf/proto" 7 | "google.golang.org/grpc/codes" 8 | ) 9 | 10 | func marshalMetadata(metadata *datacatalog.Metadata) ([]byte, error) { 11 | // if it is nil, marshal empty protobuf 12 | if metadata == nil { 13 | metadata = &datacatalog.Metadata{} 14 | } 15 | return proto.Marshal(metadata) 16 | } 17 | 18 | func unmarshalMetadata(serializedMetadata []byte) (*datacatalog.Metadata, error) { 19 | if serializedMetadata == nil { 20 | return nil, errors.NewDataCatalogErrorf(codes.Unknown, "Serialized metadata should never be nil") 21 | } 22 | var metadata datacatalog.Metadata 23 | err := proto.Unmarshal(serializedMetadata, &metadata) 24 | return &metadata, err 25 | } 26 | -------------------------------------------------------------------------------- /pkg/repositories/transformers/util_test.go: -------------------------------------------------------------------------------- 1 | package transformers 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestMarshaling(t *testing.T) { 10 | marshaledMetadata, err := marshalMetadata(metadata) 11 | assert.NoError(t, err) 12 | 13 | unmarshaledMetadata, err := unmarshalMetadata(marshaledMetadata) 14 | assert.NoError(t, err) 15 | assert.EqualValues(t, unmarshaledMetadata.KeyMap, metadata.KeyMap) 16 | } 17 | 18 | func TestMarshalingWithNil(t *testing.T) { 19 | marshaledMetadata, err := marshalMetadata(nil) 20 | assert.NoError(t, err) 21 | var expectedKeymap map[string]string 22 | unmarshaledMetadata, err := unmarshalMetadata(marshaledMetadata) 23 | assert.NoError(t, err) 24 | assert.EqualValues(t, expectedKeymap, unmarshaledMetadata.KeyMap) 25 | } 26 | -------------------------------------------------------------------------------- /pkg/repositories/utils/test_utils.go: -------------------------------------------------------------------------------- 1 | // Shared utils for postgresql tests. 2 | package utils 3 | 4 | import ( 5 | "testing" 6 | 7 | mocket "github.com/Selvatico/go-mocket" 8 | "gorm.io/driver/postgres" 9 | "gorm.io/gorm" 10 | ) 11 | 12 | func GetDbForTest(t *testing.T) *gorm.DB { 13 | mocket.Catcher.Register() 14 | db, err := gorm.Open(postgres.New(postgres.Config{DriverName: mocket.DriverName})) 15 | if err != nil { 16 | t.Fatalf("Failed to open mock db with err %v", err) 17 | } 18 | return db 19 | } 20 | -------------------------------------------------------------------------------- /pkg/runtime/application_config_provider.go: -------------------------------------------------------------------------------- 1 | package runtime 2 | 3 | import ( 4 | "context" 5 | "io/ioutil" 6 | "os" 7 | "strings" 8 | 9 | "github.com/flyteorg/flytestdlib/config" 10 | 11 | "github.com/flyteorg/datacatalog/pkg/runtime/configs" 12 | "github.com/flyteorg/flytestdlib/database" 13 | "github.com/flyteorg/flytestdlib/logger" 14 | ) 15 | 16 | const datacatalog = "datacatalog" 17 | 18 | var datacatalogConfig = config.MustRegisterSection(datacatalog, &configs.DataCatalogConfig{}) 19 | 20 | // Defines the interface to return top-level config structs necessary to start up a datacatalog application. 21 | type ApplicationConfiguration interface { 22 | GetDbConfig() *database.DbConfig 23 | GetDataCatalogConfig() configs.DataCatalogConfig 24 | } 25 | 26 | type ApplicationConfigurationProvider struct{} 27 | 28 | func (p *ApplicationConfigurationProvider) GetDbConfig() *database.DbConfig { 29 | dbConfigSection := database.GetConfig() 30 | if len(dbConfigSection.Postgres.PasswordPath) > 0 { 31 | if _, err := os.Stat(dbConfigSection.Postgres.PasswordPath); os.IsNotExist(err) { 32 | logger.Fatalf(context.Background(), 33 | "missing database password at specified path [%s]", dbConfigSection.Postgres.PasswordPath) 34 | } 35 | passwordVal, err := ioutil.ReadFile(dbConfigSection.Postgres.PasswordPath) 36 | if err != nil { 37 | logger.Fatalf(context.Background(), "failed to read database password from path [%s] with err: %v", 38 | dbConfigSection.Postgres.PasswordPath, err) 39 | } 40 | // Passwords can contain special characters as long as they are percent encoded 41 | // https://www.postgresql.org/docs/current/libpq-connect.html 42 | dbConfigSection.Postgres.Password = strings.TrimSpace(string(passwordVal)) 43 | } 44 | 45 | return dbConfigSection 46 | } 47 | 48 | func (p *ApplicationConfigurationProvider) GetDataCatalogConfig() configs.DataCatalogConfig { 49 | return *datacatalogConfig.GetConfig().(*configs.DataCatalogConfig) 50 | } 51 | 52 | func NewApplicationConfigurationProvider() ApplicationConfiguration { 53 | return &ApplicationConfigurationProvider{} 54 | } 55 | -------------------------------------------------------------------------------- /pkg/runtime/configs/data_catalog_config.go: -------------------------------------------------------------------------------- 1 | package configs 2 | 3 | import ( 4 | "time" 5 | 6 | "github.com/flyteorg/flytestdlib/config" 7 | ) 8 | 9 | //go:generate pflags DataCatalogConfig --default-var=defaultConfig 10 | 11 | var defaultConfig = &DataCatalogConfig{ 12 | StoragePrefix: "metadata", 13 | MetricsScope: "datacatalog", 14 | ProfilerPort: 10254, 15 | HeartbeatGracePeriodMultiplier: 3, 16 | MaxReservationHeartbeat: config.Duration{Duration: time.Second * 10}, 17 | } 18 | 19 | // DataCatalogConfig is the base configuration to start datacatalog 20 | type DataCatalogConfig struct { 21 | StoragePrefix string `json:"storage-prefix" pflag:",StoragePrefix specifies the prefix where DataCatalog stores offloaded ArtifactData in CloudStorage. If not specified, the data will be stored in the base container directly."` 22 | MetricsScope string `json:"metrics-scope" pflag:",Scope that the metrics will record under."` 23 | ProfilerPort int `json:"profiler-port" pflag:",Port that the profiling service is listening on."` 24 | HeartbeatGracePeriodMultiplier int `json:"heartbeat-grace-period-multiplier" pflag:",Number of heartbeats before a reservation expires without an extension."` 25 | MaxReservationHeartbeat config.Duration `json:"max-reservation-heartbeat" pflag:",The maximum available reservation extension heartbeat interval."` 26 | } 27 | -------------------------------------------------------------------------------- /pkg/runtime/configs/datacatalogconfig_flags.go: -------------------------------------------------------------------------------- 1 | // Code generated by go generate; DO NOT EDIT. 2 | // This file was generated by robots. 3 | 4 | package configs 5 | 6 | import ( 7 | "encoding/json" 8 | "reflect" 9 | 10 | "fmt" 11 | 12 | "github.com/spf13/pflag" 13 | ) 14 | 15 | // If v is a pointer, it will get its element value or the zero value of the element type. 16 | // If v is not a pointer, it will return it as is. 17 | func (DataCatalogConfig) elemValueOrNil(v interface{}) interface{} { 18 | if t := reflect.TypeOf(v); t.Kind() == reflect.Ptr { 19 | if reflect.ValueOf(v).IsNil() { 20 | return reflect.Zero(t.Elem()).Interface() 21 | } else { 22 | return reflect.ValueOf(v).Interface() 23 | } 24 | } else if v == nil { 25 | return reflect.Zero(t).Interface() 26 | } 27 | 28 | return v 29 | } 30 | 31 | func (DataCatalogConfig) mustJsonMarshal(v interface{}) string { 32 | raw, err := json.Marshal(v) 33 | if err != nil { 34 | panic(err) 35 | } 36 | 37 | return string(raw) 38 | } 39 | 40 | func (DataCatalogConfig) mustMarshalJSON(v json.Marshaler) string { 41 | raw, err := v.MarshalJSON() 42 | if err != nil { 43 | panic(err) 44 | } 45 | 46 | return string(raw) 47 | } 48 | 49 | // GetPFlagSet will return strongly types pflags for all fields in DataCatalogConfig and its nested types. The format of the 50 | // flags is json-name.json-sub-name... etc. 51 | func (cfg DataCatalogConfig) GetPFlagSet(prefix string) *pflag.FlagSet { 52 | cmdFlags := pflag.NewFlagSet("DataCatalogConfig", pflag.ExitOnError) 53 | cmdFlags.String(fmt.Sprintf("%v%v", prefix, "storage-prefix"), defaultConfig.StoragePrefix, "StoragePrefix specifies the prefix where DataCatalog stores offloaded ArtifactData in CloudStorage. If not specified, the data will be stored in the base container directly.") 54 | cmdFlags.String(fmt.Sprintf("%v%v", prefix, "metrics-scope"), defaultConfig.MetricsScope, "Scope that the metrics will record under.") 55 | cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "profiler-port"), defaultConfig.ProfilerPort, "Port that the profiling service is listening on.") 56 | cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "heartbeat-grace-period-multiplier"), defaultConfig.HeartbeatGracePeriodMultiplier, "Number of heartbeats before a reservation expires without an extension.") 57 | cmdFlags.String(fmt.Sprintf("%v%v", prefix, "max-reservation-heartbeat"), defaultConfig.MaxReservationHeartbeat.String(), "The maximum available reservation extension heartbeat interval.") 58 | return cmdFlags 59 | } 60 | -------------------------------------------------------------------------------- /pkg/runtime/configuration_provider.go: -------------------------------------------------------------------------------- 1 | package runtime 2 | 3 | // Interface for getting parsed values from a configuration file 4 | type Configuration interface { 5 | ApplicationConfiguration() ApplicationConfiguration 6 | } 7 | 8 | // Implementation of a Configuration 9 | type ConfigurationProvider struct { 10 | applicationConfiguration ApplicationConfiguration 11 | } 12 | 13 | func (p *ConfigurationProvider) ApplicationConfiguration() ApplicationConfiguration { 14 | return p.applicationConfiguration 15 | } 16 | 17 | func NewConfigurationProvider() Configuration { 18 | return &ConfigurationProvider{ 19 | applicationConfiguration: NewApplicationConfigurationProvider(), 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /pull_request_template.md: -------------------------------------------------------------------------------- 1 | ## _Read then delete this section_ 2 | 3 | _- Make sure to use a concise title for the pull-request._ 4 | 5 | _- Use #patch, #minor or #major in the pull-request title to bump the corresponding version. Otherwise, the patch version 6 | will be bumped. [More details](https://github.com/marketplace/actions/github-tag-bump)_ 7 | 8 | # TL;DR 9 | _Please replace this text with a description of what this PR accomplishes._ 10 | 11 | ## Type 12 | - [ ] Bug Fix 13 | - [ ] Feature 14 | - [ ] Plugin 15 | 16 | ## Are all requirements met? 17 | 18 | - [ ] Code completed 19 | - [ ] Smoke tested 20 | - [ ] Unit tests added 21 | - [ ] Code documentation added 22 | - [ ] Any pending items have an associated Issue 23 | 24 | ## Complete description 25 | _How did you fix the bug, make the feature etc. Link to any design docs etc_ 26 | 27 | ## Tracking Issue 28 | _Remove the '*fixes*' keyword if there will be multiple PRs to fix the linked issue_ 29 | 30 | fixes https://github.com/flyteorg/flyte/issues/ 31 | 32 | ## Follow-up issue 33 | _NA_ 34 | OR 35 | _https://github.com/flyteorg/flyte/issues/_ 36 | --------------------------------------------------------------------------------