diff --git a/.gitattributes b/.gitattributes index ff6c194874c..a0171e05ac9 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ *_pb2.py* linguist-generated +*_pb2_grpc.py* linguist-generated diff --git a/.github/workflows/check-pr-title.yml b/.github/workflows/check-pr-title.yml new file mode 100644 index 00000000000..647aaad077e --- /dev/null +++ b/.github/workflows/check-pr-title.yml @@ -0,0 +1,28 @@ +name: Check PR Title + +on: + pull_request: + branches: + - main + - '**' + +jobs: + check-title: + runs-on: ubuntu-latest + steps: + - name: Check PR Title + uses: Slashgear/action-check-pr-title@v4.3.0 + with: + regexp: '\[(ENH|BUG|DOC|TST|BLD|PERF|TYP|CLN|CHORE|RELEASE)\].*' + helpMessage: "Please tag your PR title. See https://docs.trychroma.com/contributing#contributing-code-and-ideas. You must push new code to this PR for this check to run again." + - name: Comment explaining failure + if: failure() + uses: actions/github-script@v6 + with: + script: | + github.rest.issues.createComment({ + issue_number: context.issue.number, + owner: context.repo.owner, + repo: context.repo.repo, + body: 'Please tag your PR title with one of: \\[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas' + }) diff --git a/.github/workflows/chroma-client-integration-test.yml b/.github/workflows/chroma-client-integration-test.yml index 5724959c254..e525f3a7078 100644 --- a/.github/workflows/chroma-client-integration-test.yml +++ b/.github/workflows/chroma-client-integration-test.yml @@ -15,7 +15,7 @@ jobs: timeout-minutes: 90 strategy: matrix: - python: ['3.8', '3.9', '3.10', '3.11'] + python: ['3.8', '3.9', '3.10', '3.11', '3.12'] platform: [ubuntu-latest, windows-latest] runs-on: ${{ matrix.platform }} steps: diff --git a/.github/workflows/chroma-cluster-test.yml b/.github/workflows/chroma-cluster-test.yml index e474f43ca7d..285589224ac 100644 --- a/.github/workflows/chroma-cluster-test.yml +++ b/.github/workflows/chroma-cluster-test.yml @@ -11,14 +11,15 @@ on: workflow_dispatch: jobs: - test: + test-python: strategy: matrix: python: ['3.8'] platform: ['16core-64gb-ubuntu-latest'] - testfile: ["chromadb/test/ingest/test_producer_consumer.py", - "chromadb/test/db/test_system.py", - "chromadb/test/segment/distributed/test_memberlist_provider.py",] + testfile: ["chromadb/test/db/test_system.py", + "chromadb/test/ingest/test_producer_consumer.py", + "chromadb/test/segment/distributed/test_memberlist_provider.py", + "chromadb/test/test_logservice.py"] runs-on: ${{ matrix.platform }} steps: - name: Checkout @@ -29,14 +30,40 @@ jobs: python-version: ${{ matrix.python }} - name: Install test dependencies run: python -m pip install -r requirements.txt && python -m pip install -r requirements_dev.txt - - name: Start minikube - id: minikube - uses: medyagh/setup-minikube@latest - with: - minikube-version: latest - kubernetes-version: latest - driver: docker - addons: ingress, ingress-dns - start-args: '--profile chroma-test' - - name: Integration Test - run: bin/cluster-test.sh ${{ matrix.testfile }} + - name: Install Tilt + run: | + TILT_VERSION="0.33.3" + curl -fsSL https://github.com/tilt-dev/tilt/releases/download/v$TILT_VERSION/tilt.$TILT_VERSION.linux.x86_64.tar.gz | \ + tar -xzv -C /usr/local/bin tilt + - name: Install ctlptlc + run: | + CTLPTL_VERSION="0.8.20" + curl -fsSL https://github.com/tilt-dev/ctlptl/releases/download/v$CTLPTL_VERSION/ctlptl.$CTLPTL_VERSION.linux.x86_64.tar.gz | \ + tar -xzv -C /usr/local/bin ctlptl + - name: Set up kind + run: ctlptl create cluster kind --registry=ctlptl-registry + - name: Start Tilt + run: tilt ci + - name: Test + run: bin/cluster-test.sh bash -c 'python -m pytest "${{ matrix.testfile }}"' + test-go: + runs-on: '16core-64gb-ubuntu-latest' + steps: + - name: Checkout + uses: actions/checkout@v3 + - name: Install Tilt + run: | + TILT_VERSION="0.33.3" + curl -fsSL https://github.com/tilt-dev/tilt/releases/download/v$TILT_VERSION/tilt.$TILT_VERSION.linux.x86_64.tar.gz | \ + tar -xzv -C /usr/local/bin tilt + - name: Install ctlptlc + run: | + CTLPTL_VERSION="0.8.20" + curl -fsSL https://github.com/tilt-dev/ctlptl/releases/download/v$CTLPTL_VERSION/ctlptl.$CTLPTL_VERSION.linux.x86_64.tar.gz | \ + tar -xzv -C /usr/local/bin ctlptl + - name: Set up kind + run: ctlptl create cluster kind --registry=ctlptl-registry + - name: Start Tilt + run: tilt ci + - name: Test + run: bin/cluster-test.sh bash -c 'cd go && go test -timeout 30s -run ^TestNodeWatcher$ github.com/chroma-core/chroma/go/pkg/memberlist_manager' diff --git a/.github/workflows/chroma-coordinator-test.yaml b/.github/workflows/chroma-coordinator-test.yaml index 629a9dfb146..2728f2ba5eb 100644 --- a/.github/workflows/chroma-coordinator-test.yaml +++ b/.github/workflows/chroma-coordinator-test.yaml @@ -16,8 +16,25 @@ jobs: matrix: platform: [ubuntu-latest] runs-on: ${{ matrix.platform }} + services: + postgres: + image: postgres + env: + POSTGRES_USER: chroma + POSTGRES_PASSWORD: chroma + POSTGRES_DB: chroma + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + ports: + - 5432:5432 steps: - name: Checkout uses: actions/checkout@v3 - name: Build and test - run: cd go/coordinator && make test + run: cd go && make test + env: + POSTGRES_HOST: localhost + POSTGRES_PORT: 5432 diff --git a/.github/workflows/chroma-release-python-client.yml b/.github/workflows/chroma-release-python-client.yml index 2abc0d524ab..c4f2a2990a9 100644 --- a/.github/workflows/chroma-release-python-client.yml +++ b/.github/workflows/chroma-release-python-client.yml @@ -33,7 +33,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v4 with: - python-version: '3.10' + python-version: '3.12' - name: Install Client Dev Dependencies run: python -m pip install -r ./clients/python/requirements.txt && python -m pip install -r ./clients/python/requirements_dev.txt - name: Build Client diff --git a/.github/workflows/chroma-release.yml b/.github/workflows/chroma-release.yml index 6c2250a0fb3..8fe2e8eec32 100644 --- a/.github/workflows/chroma-release.yml +++ b/.github/workflows/chroma-release.yml @@ -141,39 +141,3 @@ jobs: artifacts: "dist/chroma-${{steps.version.outputs.version}}.tar.gz" allowUpdates: true prerelease: true - - name: Trigger Hosted Chroma FE Release - uses: actions/github-script@v6 - with: - github-token: ${{ secrets.HOSTED_CHROMA_WORKFLOW_DISPATCH_TOKEN }} - script: | - const result = await github.rest.actions.createWorkflowDispatch({ - owner: 'chroma-core', - repo: 'hosted-chroma', - workflow_id: 'build-and-publish-frontend.yaml', - ref: 'main' - }) - console.log(result) - - name: Trigger Hosted Chroma Coordinator Release - uses: actions/github-script@v6 - with: - github-token: ${{ secrets.HOSTED_CHROMA_WORKFLOW_DISPATCH_TOKEN }} - script: | - const result = await github.rest.actions.createWorkflowDispatch({ - owner: 'chroma-core', - repo: 'hosted-chroma', - workflow_id: 'build-and-deploy-coordinator.yaml', - ref: 'main' - }) - console.log(result) - - name: Trigger Hosted Worker Release - uses: actions/github-script@v6 - with: - github-token: ${{ secrets.HOSTED_CHROMA_WORKFLOW_DISPATCH_TOKEN }} - script: | - const result = await github.rest.actions.createWorkflowDispatch({ - owner: 'chroma-core', - repo: 'hosted-chroma', - workflow_id: 'build-and-deploy-worker.yaml', - ref: 'main' - }) - console.log(result) diff --git a/.github/workflows/chroma-test.yml b/.github/workflows/chroma-test.yml index 12a5de4b6ed..14dc63624e9 100644 --- a/.github/workflows/chroma-test.yml +++ b/.github/workflows/chroma-test.yml @@ -16,7 +16,7 @@ jobs: timeout-minutes: 90 strategy: matrix: - python: ['3.8', '3.9', '3.10', '3.11'] + python: ['3.8', '3.9', '3.10', '3.11', '3.12'] platform: [ubuntu-latest, windows-latest] testfile: ["--ignore-glob 'chromadb/test/property/*' --ignore-glob 'chromadb/test/stress/*' --ignore='chromadb/test/auth/test_simple_rbac_authz.py'", "chromadb/test/auth/test_simple_rbac_authz.py", diff --git a/.github/workflows/release-helm-chart.yml b/.github/workflows/release-helm-chart.yml new file mode 100644 index 00000000000..a53ec02f95f --- /dev/null +++ b/.github/workflows/release-helm-chart.yml @@ -0,0 +1,30 @@ +# TODO Once distributed chroma is operational, update this to update the +# OSS helm chart we'll host. For now, just kick off the job which updates +# hosted-chroma's version. + +name: Release Helm Chart + +on: + push: + branches: + - main + paths: + - 'k8s/**' + workflow_dispatch: + +jobs: + release: + runs-on: ubuntu-latest + steps: + - name: Trigger Hosted Chroma Coordinator Release + uses: actions/github-script@v6 + with: + github-token: ${{ secrets.HOSTED_CHROMA_WORKFLOW_DISPATCH_TOKEN }} + script: | + const result = await github.rest.actions.createWorkflowDispatch({ + owner: 'chroma-core', + repo: 'hosted-chroma', + workflow_id: 'copy-oss-helm.yaml', + ref: 'main' + }) + console.log(result) \ No newline at end of file diff --git a/.gitignore b/.gitignore index 10e6fd2d731..db4dbad796a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,8 +3,9 @@ **/__pycache__ +go/bin/ +go/**/testdata/ go/coordinator/bin/ -go/coordinator/**/testdata/ *.log @@ -32,3 +33,8 @@ dist terraform.tfstate .hypothesis/ .idea + +target/ + +# environment file generated by the Javascript tests +.chroma_env diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7b1d50ec940..eb7c7916b66 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -34,3 +34,10 @@ repos: - id: mypy args: [--strict, --ignore-missing-imports, --follow-imports=silent, --disable-error-code=type-abstract, --config-file=./pyproject.toml] additional_dependencies: ["types-requests", "pydantic", "overrides", "hypothesis", "pytest", "pypika", "numpy", "types-protobuf", "kubernetes"] + + + - repo: https://github.com/pre-commit/mirrors-prettier + rev: "v3.1.0" + hooks: + - id: prettier + files: "^clients/js/.+" diff --git a/.vscode/settings.json b/.vscode/settings.json index ccddc8d4c8c..a5def08ba63 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -128,4 +128,4 @@ "unordered_set": "cpp", "algorithm": "cpp" }, -} +} \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 932b41154ab..eb620f06a2b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,11 +19,12 @@ checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" [[package]] name = "ahash" -version = "0.8.6" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91429305e9f0a25f6205c5b8e0d2db09e0708a7a6df0f42212bb56c32c8ac97a" +checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" dependencies = [ "cfg-if", + "const-random", "getrandom", "once_cell", "version_check", @@ -62,9 +63,227 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.75" +version = "1.0.81" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0952808a6c2afd1aa8947271f3a60f1a6763c7b912d210184c5149b5cf147247" + +[[package]] +name = "arc-swap" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b3d0060af21e8d11a926981cc00c6c1541aa91dd64b9f881985c3da1094425f" + +[[package]] +name = "arrow" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa285343fba4d829d49985bdc541e3789cf6000ed0e84be7c039438df4a4e78c" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-csv", + "arrow-data", + "arrow-ipc", + "arrow-json", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "753abd0a5290c1bcade7c6623a556f7d1659c5f4148b140b5b63ce7bd1a45705" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "num", +] + +[[package]] +name = "arrow-array" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d390feeb7f21b78ec997a4081a025baef1e2e0d6069e181939b61864c9779609" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown 0.14.3", + "num", +] + +[[package]] +name = "arrow-buffer" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69615b061701bcdffbc62756bc7e85c827d5290b472b580c972ebbbf690f5aa4" +dependencies = [ + "bytes", + "half", + "num", +] + +[[package]] +name = "arrow-cast" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e448e5dd2f4113bf5b74a1f26531708f5edcacc77335b7066f9398f4bcf4cdef" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "base64 0.21.7", + "chrono", + "half", + "lexical-core", + "num", +] + +[[package]] +name = "arrow-csv" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46af72211f0712612f5b18325530b9ad1bfbdc87290d5fbfd32a7da128983781" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "chrono", + "csv", + "csv-core", + "lazy_static", + "lexical-core", + "regex", +] + +[[package]] +name = "arrow-data" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67d644b91a162f3ad3135ce1184d0a31c28b816a581e08f29e8e9277a574c64e" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num", +] + +[[package]] +name = "arrow-ipc" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03dea5e79b48de6c2e04f03f62b0afea7105be7b77d134f6c5414868feefb80d" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8950719280397a47d37ac01492e3506a8a724b3fb81001900b866637a829ee0f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "indexmap 2.2.5", + "lexical-core", + "num", + "serde", + "serde_json", +] + +[[package]] +name = "arrow-ord" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ed9630979034077982d8e74a942b7ac228f33dd93a93b615b4d02ad60c260be" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "half", + "num", +] + +[[package]] +name = "arrow-row" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "007035e17ae09c4e8993e4cb8b5b96edf0afb927cd38e2dff27189b274d83dcf" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", + "hashbrown 0.14.3", +] + +[[package]] +name = "arrow-schema" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ff3e9c01f7cd169379d269f926892d0e622a704960350d09d331be3ec9e0029" + +[[package]] +name = "arrow-select" +version = "50.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4668cab20f66d8d020e1fbc0ebe47217433c1b6c8f2040faf858554e394ace6" +checksum = "1ce20973c1912de6514348e064829e50947e35977bb9d7fb637dc99ea9ffd78c" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num", +] + +[[package]] +name = "arrow-string" +version = "50.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00f3b37f2aeece31a2636d1b037dabb69ef590e03bdc7eb68519b51ec86932a7" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "num", + "regex", + "regex-syntax 0.8.2", +] [[package]] name = "async-attributes" @@ -89,13 +308,13 @@ dependencies = [ [[package]] name = "async-channel" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ca33f4bc4ed1babef42cad36cc1f51fa88be00420404e5b1e80ab1b18f7678c" +checksum = "f28243a43d821d11341ab73c80bed182dc015c514b951616cf79bd4af39af0c3" dependencies = [ "concurrent-queue", - "event-listener 4.0.0", - "event-listener-strategy", + "event-listener 5.2.0", + "event-listener-strategy 0.5.0", "futures-core", "pin-project-lite", ] @@ -106,11 +325,11 @@ version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "17ae5ebefcc48e7452b4987947920dac9450be1110cadf34d1b8c116bdbaf97c" dependencies = [ - "async-lock 3.2.0", + "async-lock 3.3.0", "async-task", "concurrent-queue", "fastrand 2.0.1", - "futures-lite 2.1.0", + "futures-lite 2.2.0", "slab", ] @@ -120,12 +339,12 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05b1b633a2115cd122d73b955eadd9916c18c8f510ec9cd1686404c60ad1c29c" dependencies = [ - "async-channel 2.1.1", + "async-channel 2.2.0", "async-executor", - "async-io 2.2.2", - "async-lock 3.2.0", + "async-io 2.3.2", + "async-lock 3.3.0", "blocking", - "futures-lite 2.1.0", + "futures-lite 2.2.0", "once_cell", ] @@ -151,18 +370,18 @@ dependencies = [ [[package]] name = "async-io" -version = "2.2.2" +version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6afaa937395a620e33dc6a742c593c01aced20aa376ffb0f628121198578ccc7" +checksum = "dcccb0f599cfa2f8ace422d3555572f47424da5648a4382a9dd0310ff8210884" dependencies = [ - "async-lock 3.2.0", + "async-lock 3.3.0", "cfg-if", "concurrent-queue", "futures-io", - "futures-lite 2.1.0", + "futures-lite 2.2.0", "parking", - "polling 3.3.1", - "rustix 0.38.28", + "polling 3.5.0", + "rustix 0.38.31", "slab", "tracing", "windows-sys 0.52.0", @@ -179,12 +398,12 @@ dependencies = [ [[package]] name = "async-lock" -version = "3.2.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7125e42787d53db9dd54261812ef17e937c95a51e4d291373b670342fa44310c" +checksum = "d034b430882f8381900d3fe6f0aaa3ad94f2cb4ac519b429692a1bc2dda4ae7b" dependencies = [ - "event-listener 4.0.0", - "event-listener-strategy", + "event-listener 4.0.3", + "event-listener-strategy 0.4.0", "pin-project-lite", ] @@ -213,7 +432,7 @@ dependencies = [ "cfg-if", "event-listener 3.1.0", "futures-lite 1.13.0", - "rustix 0.38.28", + "rustix 0.38.31", "windows-sys 0.48.0", ] @@ -223,13 +442,13 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e47d90f65a225c4527103a8d747001fc56e375203592b25ad103e1ca13124c5" dependencies = [ - "async-io 2.2.2", + "async-io 2.3.2", "async-lock 2.8.0", "atomic-waker", "cfg-if", "futures-core", "futures-io", - "rustix 0.38.28", + "rustix 0.38.31", "signal-hook-registry", "slab", "windows-sys 0.48.0", @@ -282,24 +501,24 @@ checksum = "16e62a023e7c117e27523144c5d2459f4397fcc3cab0085af8e2224f643a0193" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] name = "async-task" -version = "4.5.0" +version = "4.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4eb2cdb97421e01129ccb49169d8279ed21e829929144f4a22a6e54ac549ca1" +checksum = "fbb36e985947064623dbd357f727af08ffd077f93d696782f3c56365fa2e2799" [[package]] name = "async-trait" -version = "0.1.74" +version = "0.1.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a66537f1bb974b254c98ed142ff995236e81b9d0fe4db0575f46612cb15eb0f9" +checksum = "c980ee35e870bd1a4d2c8294d4c04d0499e67bca1e4b5cefcc693c2fa00caea9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -338,12 +557,11 @@ checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" [[package]] name = "aws-config" -version = "1.1.2" +version = "1.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e64b72d4bdbb41a73d27709c65a25b6e4bfc8321bf70fa3a8b19ce7d4eb81b0" +checksum = "0b96342ea8948ab9bef3e6234ea97fc32e2d8a88d8fb6a084e52267317f94b6b" dependencies = [ "aws-credential-types", - "aws-http", "aws-runtime", "aws-sdk-sso", "aws-sdk-ssooidc", @@ -358,7 +576,7 @@ dependencies = [ "bytes", "fastrand 2.0.1", "hex", - "http", + "http 0.2.12", "hyper", "ring", "time", @@ -369,9 +587,9 @@ dependencies = [ [[package]] name = "aws-credential-types" -version = "1.1.2" +version = "1.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a7cb3510b95492bd9014b60e2e3bee3e48bc516e220316f8e6b60df18b47331" +checksum = "273fa47dafc9ef14c2c074ddddbea4561ff01b7f68d5091c0e9737ced605c01d" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -379,30 +597,13 @@ dependencies = [ "zeroize", ] -[[package]] -name = "aws-http" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a95d41abe4e941399fdb4bc2f54713eac3c839d98151875948bb24e66ab658f2" -dependencies = [ - "aws-smithy-runtime-api", - "aws-smithy-types", - "aws-types", - "bytes", - "http", - "http-body", - "pin-project-lite", - "tracing", -] - [[package]] name = "aws-runtime" -version = "1.1.2" +version = "1.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "233cca219c6705d525ace011d6f9bc51aaf32fce5b4c41661d2d7ff22d9b4d49" +checksum = "6e38bab716c8bf07da24be07ecc02e0f5656ce8f30a891322ecdcb202f943b85" dependencies = [ "aws-credential-types", - "aws-http", "aws-sigv4", "aws-smithy-async", "aws-smithy-eventstream", @@ -410,21 +611,23 @@ dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", "aws-types", + "bytes", "fastrand 2.0.1", - "http", + "http 0.2.12", + "http-body", "percent-encoding", + "pin-project-lite", "tracing", "uuid", ] [[package]] name = "aws-sdk-s3" -version = "1.12.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "634fbe5b6591ee2e281cd2ba8641e9bd752dbf5bf338924d6ad4bd5a3304fe31" +checksum = "93d35d39379445970fc3e4ddf7559fff2c32935ce0b279f9cb27080d6b7c6d94" dependencies = [ "aws-credential-types", - "aws-http", "aws-runtime", "aws-sigv4", "aws-smithy-async", @@ -438,7 +641,7 @@ dependencies = [ "aws-smithy-xml", "aws-types", "bytes", - "http", + "http 0.2.12", "http-body", "once_cell", "percent-encoding", @@ -449,12 +652,11 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.10.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee41005e0f3a19ae749c7953d9e1f1ef8d2183f76f64966e346fa41c1ba0ed44" +checksum = "d84bd3925a17c9adbf6ec65d52104a44a09629d8f70290542beeee69a95aee7f" dependencies = [ "aws-credential-types", - "aws-http", "aws-runtime", "aws-smithy-async", "aws-smithy-http", @@ -464,7 +666,7 @@ dependencies = [ "aws-smithy-types", "aws-types", "bytes", - "http", + "http 0.2.12", "once_cell", "regex-lite", "tracing", @@ -472,12 +674,11 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.10.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa08168f8a27505e7b90f922c32a489feb1f2133878981a15138bebc849ac09c" +checksum = "2c2dae39e997f58bc4d6292e6244b26ba630c01ab671b6f9f44309de3eb80ab8" dependencies = [ "aws-credential-types", - "aws-http", "aws-runtime", "aws-smithy-async", "aws-smithy-http", @@ -487,7 +688,7 @@ dependencies = [ "aws-smithy-types", "aws-types", "bytes", - "http", + "http 0.2.12", "once_cell", "regex-lite", "tracing", @@ -495,12 +696,11 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.10.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29102eff04d50ef70f11a48823db33e33c6cc5f027bfb6ff4864efbd5f1f66f3" +checksum = "17fd9a53869fee17cea77e352084e1aa71e2c5e323d974c13a9c2bcfd9544c7f" dependencies = [ "aws-credential-types", - "aws-http", "aws-runtime", "aws-smithy-async", "aws-smithy-http", @@ -511,7 +711,7 @@ dependencies = [ "aws-smithy-types", "aws-smithy-xml", "aws-types", - "http", + "http 0.2.12", "once_cell", "regex-lite", "tracing", @@ -519,9 +719,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.1.2" +version = "1.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b92384b39aedb258aa734fe0e7b2ffcd13f33e68227251a72cd2635e0acc8f1a" +checksum = "8ada00a4645d7d89f296fe0ddbc3fe3554f03035937c849a05d37ddffc1f29a1" dependencies = [ "aws-credential-types", "aws-smithy-eventstream", @@ -533,7 +733,8 @@ dependencies = [ "form_urlencoded", "hex", "hmac", - "http", + "http 0.2.12", + "http 1.1.0", "once_cell", "p256 0.11.1", "percent-encoding", @@ -547,9 +748,9 @@ dependencies = [ [[package]] name = "aws-smithy-async" -version = "1.1.2" +version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71d8e1c0904f78c76846a9dad41c28b41d330d97741c3e70d003d9a747d95e2a" +checksum = "d26ea8fa03025b2face2b3038a63525a10891e3d8829901d502e5384a0d8cd46" dependencies = [ "futures-util", "pin-project-lite", @@ -558,9 +759,9 @@ dependencies = [ [[package]] name = "aws-smithy-checksums" -version = "0.60.2" +version = "0.60.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62d59ef74bf94562512e570eeccb81e9b3879f9136b2171ed4bf996ffa609955" +checksum = "83fa43bc04a6b2441968faeab56e68da3812f978a670a5db32accbdcafddd12f" dependencies = [ "aws-smithy-http", "aws-smithy-types", @@ -568,7 +769,7 @@ dependencies = [ "crc32c", "crc32fast", "hex", - "http", + "http 0.2.12", "http-body", "md-5", "pin-project-lite", @@ -579,9 +780,9 @@ dependencies = [ [[package]] name = "aws-smithy-eventstream" -version = "0.60.2" +version = "0.60.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31cf0466890a20988b9b2864250dd907f769bd189af1a51ba67beec86f7669fb" +checksum = "e6363078f927f612b970edf9d1903ef5cef9a64d1e8423525ebb1f0a1633c858" dependencies = [ "aws-smithy-types", "bytes", @@ -590,9 +791,9 @@ dependencies = [ [[package]] name = "aws-smithy-http" -version = "0.60.2" +version = "0.60.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "568a3b159001358dd96143378afd7470e19baffb6918e4b5016abe576e553f9c" +checksum = "3f10fa66956f01540051b0aa7ad54574640f748f9839e843442d99b970d3aff9" dependencies = [ "aws-smithy-eventstream", "aws-smithy-runtime-api", @@ -600,7 +801,7 @@ dependencies = [ "bytes", "bytes-utils", "futures-core", - "http", + "http 0.2.12", "http-body", "once_cell", "percent-encoding", @@ -611,18 +812,18 @@ dependencies = [ [[package]] name = "aws-smithy-json" -version = "0.60.2" +version = "0.60.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f12bfb23370a069f8facbfd53ce78213461b0a8570f6c81488030f5ab6f8cc4e" +checksum = "4683df9469ef09468dad3473d129960119a0d3593617542b7d52086c8486f2d6" dependencies = [ "aws-smithy-types", ] [[package]] name = "aws-smithy-query" -version = "0.60.2" +version = "0.60.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b1adc06e0175c175d280267bb8fd028143518013fcb869e1c3199569a2e902a" +checksum = "f2fbd61ceb3fe8a1cb7352e42689cec5335833cd9f94103a61e98f9bb61c64bb" dependencies = [ "aws-smithy-types", "urlencoding", @@ -630,9 +831,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.1.2" +version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7cf0f6845d2d97b953cea791b0ee37191c5509f2897ec7eb7580a0e7a594e98b" +checksum = "ec81002d883e5a7fd2bb063d6fb51c4999eb55d404f4fff3dd878bf4733b9f01" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -641,7 +842,7 @@ dependencies = [ "bytes", "fastrand 2.0.1", "h2", - "http", + "http 0.2.12", "http-body", "hyper", "hyper-rustls", @@ -655,14 +856,15 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.1.2" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47798ba97a33979c80e837519cf837f18fd6df0adb02dd5286a75d9891c6e671" +checksum = "9acb931e0adaf5132de878f1398d83f8677f90ba70f01f65ff87f6d7244be1c5" dependencies = [ "aws-smithy-async", "aws-smithy-types", "bytes", - "http", + "http 0.2.12", + "http 1.1.0", "pin-project-lite", "tokio", "tracing", @@ -671,15 +873,15 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.1.2" +version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e9a85eafeaf783b2408e35af599e8b96f2c49d9a5d13ad3a887fbdefb6bc744" +checksum = "abe14dceea1e70101d38fbf2a99e6a34159477c0fb95e68e05c66bd7ae4c3729" dependencies = [ "base64-simd", "bytes", "bytes-utils", "futures-core", - "http", + "http 0.2.12", "http-body", "itoa", "num-integer", @@ -694,24 +896,24 @@ dependencies = [ [[package]] name = "aws-smithy-xml" -version = "0.60.2" +version = "0.60.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a84bee2b44c22cbba59f12c34b831a97df698f8e43df579b35998652a00dc13" +checksum = "872c68cf019c0e4afc5de7753c4f7288ce4b71663212771bf5e4542eb9346ca9" dependencies = [ "xmlparser", ] [[package]] name = "aws-types" -version = "1.1.2" +version = "1.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8549aa62c5b7db5c57ab915200ee214b4f5d8f19b29a4a8fa0b3ad3bca1380e3" +checksum = "d07c63521aa1ea9a9f92a701f1a08ce3fd20b46c6efc0d5c8947c1fd879e3df1" dependencies = [ "aws-credential-types", "aws-smithy-async", "aws-smithy-runtime-api", "aws-smithy-types", - "http", + "http 0.2.12", "rustc_version", "tracing", ] @@ -727,7 +929,7 @@ dependencies = [ "bitflags 1.3.2", "bytes", "futures-util", - "http", + "http 0.2.12", "http-body", "hyper", "itoa", @@ -753,7 +955,7 @@ dependencies = [ "async-trait", "bytes", "futures-util", - "http", + "http 0.2.12", "http-body", "mime", "rustversion", @@ -807,9 +1009,9 @@ checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" [[package]] name = "base64" -version = "0.21.5" +version = "0.21.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35636a1494ede3b646cc98f74f8e62c773a38a659ebc777a2cf26b9b74171df9" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" [[package]] name = "base64-simd" @@ -827,6 +1029,15 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + [[package]] name = "bit-vec" version = "0.6.3" @@ -841,9 +1052,18 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.4.1" +version = "2.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "327762f6e5a765692301e5bb513e0d9fef63be86bbc14528052b1cd3e6f03e07" +checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf" + +[[package]] +name = "bitpacking" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8c7d2ac73c167c06af4a5f37e6e59d84148d57ccbe4480b76f0273eefea82d7" +dependencies = [ + "crunchy", +] [[package]] name = "block-buffer" @@ -860,27 +1080,27 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a37913e8dc4ddcc604f0c6d3bf2887c995153af3611de9e23c352b44c1b9118" dependencies = [ - "async-channel 2.1.1", - "async-lock 3.2.0", + "async-channel 2.2.0", + "async-lock 3.3.0", "async-task", "fastrand 2.0.1", "futures-io", - "futures-lite 2.1.0", + "futures-lite 2.2.0", "piper", "tracing", ] [[package]] name = "bumpalo" -version = "3.14.0" +version = "3.15.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f30e7476521f6f8af1a1c4c0b8cc94f0bee37d91763d0ca2665f299b6cd8aec" +checksum = "7ff69b9dd49fd426c69a0db9fc04dd934cdb6645ff000864d98f7e2af8830eaa" [[package]] name = "bytemuck" -version = "1.14.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "374d28ec25809ee0e23827c2ab573d729e293f281dfe393500e7ad618baa61c6" +checksum = "5d6d68c57235a3a081186990eca2867354726650f42f7516ca50c28d6281fd15" [[package]] name = "byteorder" @@ -906,14 +1126,20 @@ dependencies = [ [[package]] name = "cc" -version = "1.0.83" +version = "1.0.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0" +checksum = "8cd6604a82acf3039f1144f54b8eb34e91ffba622051189e71b781822d5ee1f5" dependencies = [ "jobserver", "libc", ] +[[package]] +name = "census" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f4c707c6a209cbe82d10abd08e1ea8995e9ea937d2550646e02798948992be0" + [[package]] name = "cfg-if" version = "1.0.0" @@ -922,9 +1148,9 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "chrono" -version = "0.4.31" +version = "0.4.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f2c685bad3eb3d45a01354cedb7d5faa66194d1d58ba6e267a8de788f79db38" +checksum = "8eaf5903dcbc0a39312feb77df2ff4c76387d591b9fc7b04a238dcf8bb62639a" dependencies = [ "android-tzdata", "iana-time-zone", @@ -932,7 +1158,7 @@ dependencies = [ "num-traits", "serde", "wasm-bindgen", - "windows-targets 0.48.5", + "windows-targets 0.52.4", ] [[package]] @@ -946,9 +1172,29 @@ dependencies = [ [[package]] name = "const-oid" -version = "0.9.5" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + +[[package]] +name = "const-random" +version = "0.1.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28c122c3980598d243d63d9a704629a2d748d101f278052ff068be5a4423ab6f" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom", + "once_cell", + "tiny-keccak", +] [[package]] name = "core-foundation" @@ -968,9 +1214,9 @@ checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f" [[package]] name = "cpufeatures" -version = "0.2.11" +version = "0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce420fe07aecd3e67c5f910618fe65e94158f6dcc0adf44e00d69ce2bdfe0fd0" +checksum = "53fe5e26ff1b7aef8bca9c6080520cfb8d9333c7568e1829cef191a9723e5504" dependencies = [ "libc", ] @@ -992,54 +1238,61 @@ checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" [[package]] name = "crc32c" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8f48d60e5b4d2c53d5c2b1d8a58c849a70ae5e5509b08a48d047e3b65714a74" +checksum = "89254598aa9b9fa608de44b3ae54c810f0f06d755e24c50177f1f8f31ff50ce2" dependencies = [ "rustc_version", ] [[package]] name = "crc32fast" -version = "1.3.2" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d" +checksum = "b3855a8a784b474f333699ef2bbca9db2c4a1f6d9088a90a2d25b1eb53111eaa" dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab3db02a9c5b5121e1e42fbdb1aeb65f5e02624cc58c43f2884c6ccac0b82f95" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" -version = "0.8.3" +version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" dependencies = [ - "cfg-if", "crossbeam-epoch", "crossbeam-utils", ] [[package]] name = "crossbeam-epoch" -version = "0.9.15" +version = "0.9.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" dependencies = [ - "autocfg", - "cfg-if", "crossbeam-utils", - "memoffset", - "scopeguard", ] [[package]] name = "crossbeam-utils" -version = "0.8.16" +version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294" -dependencies = [ - "cfg-if", -] +checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" + +[[package]] +name = "crunchy" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" [[package]] name = "crypto-bigint" @@ -1075,11 +1328,32 @@ dependencies = [ "typenum", ] +[[package]] +name = "csv" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac574ff4d437a7b5ad237ef331c17ccca63c46479e5b5453eb8e10bb99a759fe" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "csv-core" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70" +dependencies = [ + "memchr", +] + [[package]] name = "curve25519-dalek" -version = "4.1.1" +version = "4.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e89b8c6a2e4b1f45971ad09761aafb85514a84744b67a95e32c3cc1352d1f65c" +checksum = "0a677b8922c94e01bdbb12126b0bc852f00447528dee1782229af9c720c3f348" dependencies = [ "cfg-if", "cpufeatures", @@ -1100,14 +1374,14 @@ checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] name = "darling" -version = "0.20.3" +version = "0.20.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0209d94da627ab5605dcccf08bb18afa5009cfbef48d8a8b7d7bdbc79be25c5e" +checksum = "54e36fcd13ed84ffdfda6f5be89b31287cbb80c439841fe69e04841435464391" dependencies = [ "darling_core", "darling_macro", @@ -1115,27 +1389,27 @@ dependencies = [ [[package]] name = "darling_core" -version = "0.20.3" +version = "0.20.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "177e3443818124b357d8e76f53be906d60937f0d3a90773a664fa63fa253e621" +checksum = "9c2cf1c23a687a1feeb728783b993c4e1ad83d99f351801977dd809b48d0a70f" dependencies = [ "fnv", "ident_case", "proc-macro2", "quote", "strsim", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] name = "darling_macro" -version = "0.20.3" +version = "0.20.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836a9bbc7ad63342d6d6e7b815ccab164bc77a2d95d84bc3117a8c0d5c98e2d5" +checksum = "a668eda54683121533a393014d8692171709ff57a7d61f187b6e782719f8933f" dependencies = [ "darling_core", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -1167,9 +1441,9 @@ dependencies = [ [[package]] name = "deranged" -version = "0.3.10" +version = "0.3.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8eb30d70a07a3b04884d2677f06bec33509dc67ca60d92949e5535352d3191dc" +checksum = "b42b6fa04a440b495c8b04d0e71b707c585f83cb9cb28cf8cd0d976c315e31b4" dependencies = [ "powerfmt", "serde", @@ -1198,11 +1472,17 @@ dependencies = [ "subtle", ] +[[package]] +name = "downcast-rs" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea835d29036a4087793836fa931b08837ad5e957da9e23886b29586fb9b6650" + [[package]] name = "dyn-clone" -version = "1.0.16" +version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "545b22097d44f8a9581187cdf93de7a71e4722bf51200cfaba810865b49a495d" +checksum = "0d6ef0072f8a535281e4876be788938b528e9a1d43900b82c2569af7da799125" [[package]] name = "ecdsa" @@ -1242,9 +1522,9 @@ dependencies = [ [[package]] name = "ed25519-dalek" -version = "2.1.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f628eaec48bfd21b865dc2950cfa014450c01d2fa2b69a86c2fd5844ec523c0" +checksum = "4a3daa8e81a3963a60642bcc1f90a670680bd4a77535faa384e9d1c79d620871" dependencies = [ "curve25519-dalek", "ed25519", @@ -1256,9 +1536,9 @@ dependencies = [ [[package]] name = "either" -version = "1.9.0" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07" +checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" [[package]] name = "elliptic-curve" @@ -1345,9 +1625,20 @@ dependencies = [ [[package]] name = "event-listener" -version = "4.0.0" +version = "4.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "770d968249b5d99410d61f5bf89057f3199a077a04d087092f58e7d10692baae" +checksum = "67b215c49b2b248c855fb73579eb1f4f26c38ffdc12973e20e07b91d78d5646e" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener" +version = "5.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b5fb89194fa3cad959b833185b3063ba881dbfc7030680b314250779fb4cc91" dependencies = [ "concurrent-queue", "parking", @@ -1360,10 +1651,26 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "958e4d70b6d5e81971bebec42271ec641e7ff4e170a6fa605f2b8a8b65cb97d3" dependencies = [ - "event-listener 4.0.0", + "event-listener 4.0.3", "pin-project-lite", ] +[[package]] +name = "event-listener-strategy" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "feedafcaa9b749175d5ac357452a9d41ea2911da598fde46ce1fe02c37751291" +dependencies = [ + "event-listener 5.2.0", + "pin-project-lite", +] + +[[package]] +name = "fastdivide" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25c7df09945d65ea8d70b3321547ed414bbc540aad5bac6883d021b970f35b04" + [[package]] name = "fastrand" version = "1.9.0" @@ -1401,15 +1708,15 @@ dependencies = [ [[package]] name = "fiat-crypto" -version = "0.2.5" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27573eac26f4dd11e2b1916c3fe1baa56407c83c71a773a8ba17ec0bca03b6b7" +checksum = "1676f435fc1dadde4d03e43f5d62b259e1ce5f40bd4ffb21db2b42ebe59c1382" [[package]] name = "figment" -version = "0.10.12" +version = "0.10.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "649f3e5d826594057e9a519626304d8da859ea8a0b18ce99500c586b8d45faee" +checksum = "2b6e5bc7bd59d60d0d45a6ccab6cf0f4ce28698fb4e81e750ddf229c9b824026" dependencies = [ "atomic", "parking_lot", @@ -1427,6 +1734,16 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" +[[package]] +name = "flatbuffers" +version = "23.5.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dac53e22462d78c16d64a1cd22371b54cc3fe94aa15e7886a2fa6e5d1ab8640" +dependencies = [ + "bitflags 1.3.2", + "rustc_version", +] + [[package]] name = "flate2" version = "1.0.28" @@ -1467,11 +1784,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs4" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2eeb4ed9e12f43b7fa0baae3f9cdda28352770132ef2e09a23760c29cae8bd47" +dependencies = [ + "rustix 0.38.31", + "windows-sys 0.48.0", +] + [[package]] name = "futures" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da0290714b38af9b4a7b094b8a37086d1b4e61f2df9122c3cad2577669145335" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" dependencies = [ "futures-channel", "futures-core", @@ -1484,9 +1811,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff4dd66668b557604244583e3e1e1eada8c5c2e96a6d0d6653ede395b78bbacb" +checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" dependencies = [ "futures-core", "futures-sink", @@ -1494,15 +1821,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb1d22c66e66d9d72e1758f0bd7d4fd0bee04cad842ee34587d68c07e45d088c" +checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" [[package]] name = "futures-executor" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f4fb8693db0cf099eadcca0efe2a5a22e4550f98ed16aba6c48700da29597bc" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" dependencies = [ "futures-core", "futures-task", @@ -1511,9 +1838,9 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bf34a163b5c4c52d0478a4d757da8fb65cabef42ba90515efee0f6f9fa45aaa" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" [[package]] name = "futures-lite" @@ -1532,9 +1859,9 @@ dependencies = [ [[package]] name = "futures-lite" -version = "2.1.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aeee267a1883f7ebef3700f262d2d54de95dfaf38189015a74fdc4e0c7ad8143" +checksum = "445ba825b27408685aaecefd65178908c36c6e96aaf6d8599419d46e624192ba" dependencies = [ "fastrand 2.0.1", "futures-core", @@ -1545,38 +1872,38 @@ dependencies = [ [[package]] name = "futures-macro" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53b153fd91e4b0147f4aced87be237c98248656bb01050b96bf3ee89220a8ddb" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] name = "futures-sink" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e36d3378ee38c2a36ad710c5d30c2911d752cb941c00c72dbabfb786a7970817" +checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" [[package]] name = "futures-task" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efd193069b0ddadc69c46389b740bbccdd97203899b48d09c5f7969591d6bae2" +checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" [[package]] name = "futures-timer" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e64b03909df88034c26dc1547e8970b91f98bdb65165d6a4e9110d94263dbb2c" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" [[package]] name = "futures-util" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a19526d624e703a3179b3d322efec918b6246ea0fa51d41124525f00f1cc8104" +checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" dependencies = [ "futures-channel", "futures-core", @@ -1590,6 +1917,19 @@ dependencies = [ "slab", ] +[[package]] +name = "generator" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc16584ff22b460a382b7feec54b23d2908d858152e5739a120b949293bd74e" +dependencies = [ + "cc", + "libc", + "log", + "rustversion", + "windows", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -1603,9 +1943,9 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.2.11" +version = "0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe9006bed769170c11f845cf00c7c1e9092aeb3f268e007c3e760ac68008070f" +checksum = "190092ea657667030ac6a35e305e62fc4dd69fd98ac98631e5d3a2b1575a12b5" dependencies = [ "cfg-if", "js-sys", @@ -1656,23 +1996,34 @@ dependencies = [ [[package]] name = "h2" -version = "0.3.22" +version = "0.3.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d6250322ef6e60f93f9a2162799302cd6f68f79f6e5d85c8c16f14d1d958178" +checksum = "bb2c4422095b67ee78da96fbb51a4cc413b3b25883c7717ff7ca1ab31022c9c9" dependencies = [ "bytes", "fnv", "futures-core", "futures-sink", "futures-util", - "http", - "indexmap 2.1.0", + "http 0.2.12", + "indexmap 2.2.5", "slab", "tokio", "tokio-util", "tracing", ] +[[package]] +name = "half" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5eceaaeec696539ddaf7b333340f1af35a5aa87ae3e4f3ead0532f72affab2e" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1697,9 +2048,9 @@ checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" [[package]] name = "hermit-abi" -version = "0.3.3" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" [[package]] name = "hex" @@ -1727,18 +2078,35 @@ dependencies = [ [[package]] name = "home" -version = "0.5.5" +version = "0.5.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5444c27eef6923071f7ebcc33e3444508466a76f7a2b93da00ed6e19f30c1ddb" +checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.52.0", +] + +[[package]] +name = "htmlescape" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9025058dae765dee5070ec375f591e2ba14638c63feff74f13805a72e523163" + +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", ] [[package]] name = "http" -version = "0.2.11" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8947b1a6fad4393052c7ba1f4cd97bed3e953a95c79c92ad9b051a04611d9fbb" +checksum = "21b9ddb458710bc376481b842f5da65cdf31522de232c1ca8146abce2a358258" dependencies = [ "bytes", "fnv", @@ -1752,7 +2120,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" dependencies = [ "bytes", - "http", + "http 0.2.12", "pin-project-lite", ] @@ -1776,22 +2144,22 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hyper" -version = "0.14.27" +version = "0.14.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffb1cfd654a8219eaef89881fdb3bb3b1cdc5fa75ded05d6933b2b382e395468" +checksum = "bf96e135eb83a2a8ddf766e426a841d8ddd7449d5f00d34ea02b41d2f19eef80" dependencies = [ "bytes", "futures-channel", "futures-core", "futures-util", "h2", - "http", + "http 0.2.12", "http-body", "httparse", "httpdate", "itoa", "pin-project-lite", - "socket2 0.4.10", + "socket2 0.5.6", "tokio", "tower-service", "tracing", @@ -1805,7 +2173,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" dependencies = [ "futures-util", - "http", + "http 0.2.12", "hyper", "log", "rustls", @@ -1828,9 +2196,9 @@ dependencies = [ [[package]] name = "iana-time-zone" -version = "0.1.58" +version = "0.1.60" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8326b86b6cff230b97d0d312a6c40a60726df3332e721f72a1b035f451663b20" +checksum = "e7ffbb5a1b541ea2561f8c41c087286cc091e21e556a4f09a8f6cbf17b69b141" dependencies = [ "android_system_properties", "core-foundation-sys", @@ -1878,9 +2246,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.1.0" +version = "2.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d530e1a18b1cb4c484e6e34556a0d948706958449fca0cab753d649f2bce3d1f" +checksum = "7b0b929d511467233429c45a44ac1dcaa21ba0f5ba11e4879e6ed28ddb4f9df4" dependencies = [ "equivalent", "hashbrown 0.14.3", @@ -1900,6 +2268,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c" dependencies = [ "cfg-if", + "js-sys", + "wasm-bindgen", + "web-sys", ] [[package]] @@ -1945,18 +2316,18 @@ checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c" [[package]] name = "jobserver" -version = "0.1.27" +version = "0.1.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c37f63953c4c63420ed5fd3d6d398c719489b9f872b9fa683262f8edd363c7d" +checksum = "ab46a6e9526ddef3ae7f787c06f0f2600639ba80ea3eade3d8e670a2230f51d6" dependencies = [ "libc", ] [[package]] name = "js-sys" -version = "0.3.66" +version = "0.3.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cee9c64da59eae3b50095c18d3e74f8b73c0b86d2792824ff01bbce68ba229ca" +checksum = "29c15563dc2726973df627357ce0c9ddddbea194836909d655df6a75d2cf296d" dependencies = [ "wasm-bindgen", ] @@ -1974,14 +2345,16 @@ dependencies = [ ] [[package]] -name = "jsonpath_lib" -version = "0.3.0" +name = "jsonpath-rust" +version = "0.3.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eaa63191d68230cccb81c5aa23abd53ed64d83337cacbb25a7b8c7979523774f" +checksum = "06cc127b7c3d270be504572364f9569761a180b981919dd0d87693a7f5fb7829" dependencies = [ - "log", - "serde", + "pest", + "pest_derive", + "regex", "serde_json", + "thiserror", ] [[package]] @@ -1990,7 +2363,7 @@ version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "edc3606fd16aca7989db2f84bb25684d0270c6d6fa1dbcd0025af7b4130523a6" dependencies = [ - "base64 0.21.5", + "base64 0.21.7", "bytes", "chrono", "serde", @@ -2000,9 +2373,9 @@ dependencies = [ [[package]] name = "kube" -version = "0.87.1" +version = "0.87.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e34392aea935145070dcd5b39a6dea689ac6534d7d117461316c3d157b1d0fc3" +checksum = "3499c8d60c763246c7a213f51caac1e9033f46026904cb89bc8951ae8601f26e" dependencies = [ "k8s-openapi", "kube-client", @@ -2013,22 +2386,22 @@ dependencies = [ [[package]] name = "kube-client" -version = "0.87.1" +version = "0.87.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7266548b9269d9fa19022620d706697e64f312fb2ba31b93e6986453fcc82c92" +checksum = "033450dfa0762130565890dadf2f8835faedf749376ca13345bcd8ecd6b5f29f" dependencies = [ - "base64 0.21.5", + "base64 0.21.7", "bytes", "chrono", "either", "futures", "home", - "http", + "http 0.2.12", "http-body", "hyper", "hyper-rustls", "hyper-timeout", - "jsonpath_lib", + "jsonpath-rust", "k8s-openapi", "kube-core", "pem", @@ -2049,13 +2422,13 @@ dependencies = [ [[package]] name = "kube-core" -version = "0.87.1" +version = "0.87.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8321c315b96b59f59ef6b33f604b84b905ab8f9ff114a4f909d934c520227b1" +checksum = "b5bba93d054786eba7994d03ce522f368ef7d48c88a1826faa28478d85fb63ae" dependencies = [ "chrono", "form_urlencoded", - "http", + "http 0.2.12", "json-patch", "k8s-openapi", "once_cell", @@ -2067,22 +2440,22 @@ dependencies = [ [[package]] name = "kube-derive" -version = "0.87.1" +version = "0.87.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d54591e1f37fc329d412c0fdaced010cc1305b546a39f283fc51700f8fb49421" +checksum = "91e98dd5e5767c7b894c1f0e41fd628b145f808e981feb8b08ed66455d47f1a4" dependencies = [ "darling", "proc-macro2", "quote", "serde_json", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] name = "kube-runtime" -version = "0.87.1" +version = "0.87.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e511e2c1a368d9d4bf6e70db58197e535d818df355b5a2007a8aeb17a370a8ba" +checksum = "2d8893eb18fbf6bb6c80ef6ee7dd11ec32b1dc3c034c988ac1b3a84d46a230ae" dependencies = [ "ahash", "async-trait", @@ -2122,11 +2495,81 @@ dependencies = [ "spin 0.5.2", ] +[[package]] +name = "levenshtein_automata" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c2cdeb66e45e9f36bfad5bbdb4d2384e70936afbee843c6f6543f0c551ebb25" + +[[package]] +name = "lexical-core" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cde5de06e8d4c2faabc400238f9ae1c74d5412d03a7bd067645ccbc47070e46" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683b3a5ebd0130b8fb52ba0bdc718cc56815b6a097e28ae5a6997d0ad17dc05f" +dependencies = [ + "lexical-parse-integer", + "lexical-util", + "static_assertions", +] + +[[package]] +name = "lexical-parse-integer" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d0994485ed0c312f6d965766754ea177d07f9c00c9b82a5ee62ed5b47945ee9" +dependencies = [ + "lexical-util", + "static_assertions", +] + +[[package]] +name = "lexical-util" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5255b9ff16ff898710eb9eb63cb39248ea8a5bb036bea8085b1a767ff6c4e3fc" +dependencies = [ + "static_assertions", +] + +[[package]] +name = "lexical-write-float" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accabaa1c4581f05a3923d1b4cfd124c329352288b7b9da09e766b0668116862" +dependencies = [ + "lexical-util", + "lexical-write-integer", + "static_assertions", +] + +[[package]] +name = "lexical-write-integer" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1b6f3d1f4422866b68192d62f77bc5c700bee84f3069f2469d7bc8c77852446" +dependencies = [ + "lexical-util", + "static_assertions", +] + [[package]] name = "libc" -version = "0.2.151" +version = "0.2.153" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "302d7ab3130588088d277783b1e2d2e10c9e9e4a16dd9050e6ec93fb3e7048f4" +checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" [[package]] name = "libm" @@ -2142,9 +2585,9 @@ checksum = "ef53942eb7bf7ff43a617b3e2c1c4a5ecf5944a7c1bc12d7ee39bbb15e5c1519" [[package]] name = "linux-raw-sys" -version = "0.4.12" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4cd1a83af159aa67994778be9070f0ae1bd732942279cabb14f86f986a21456" +checksum = "01cda141df6706de531b6c46c3a33ecca755538219bd484262fa09410c13539c" [[package]] name = "lock_api" @@ -2158,13 +2601,36 @@ dependencies = [ [[package]] name = "log" -version = "0.4.20" +version = "0.4.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" +checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" dependencies = [ "value-bag", ] +[[package]] +name = "loom" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff50ecb28bb86013e935fb6683ab1f6d3a20016f123c76fd4c27470076ac30f5" +dependencies = [ + "cfg-if", + "generator", + "pin-utils", + "scoped-tls", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "lru" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4a83fb7698b3643a0e34f9ae6f2e8f0178c0fd42f8b59d493aa271ff3a5bf21" +dependencies = [ + "hashbrown 0.14.3", +] + [[package]] name = "lz4" version = "1.24.0" @@ -2185,6 +2651,21 @@ dependencies = [ "libc", ] +[[package]] +name = "lz4_flex" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "912b45c753ff5f7f5208307e8ace7d2a2e30d024e26d3509f3dce546c044ce15" + +[[package]] +name = "matchers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +dependencies = [ + "regex-automata 0.1.10", +] + [[package]] name = "matchit" version = "0.7.3" @@ -2201,19 +2682,29 @@ dependencies = [ "digest", ] +[[package]] +name = "measure_time" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56220900f1a0923789ecd6bf25fbae8af3b2f1ff3e9e297fc9b6b8674dd4d852" +dependencies = [ + "instant", + "log", +] + [[package]] name = "memchr" -version = "2.6.4" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167" +checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149" [[package]] -name = "memoffset" -version = "0.9.0" +name = "memmap2" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c" +checksum = "f49388d20533534cd19360ad3d6a7dadc885944aa802ba3995040c5ec11288c6" dependencies = [ - "autocfg", + "libc", ] [[package]] @@ -2230,18 +2721,18 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "miniz_oxide" -version = "0.7.1" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7" +checksum = "9d811f3e15f28568be3407c8e7fdb6514c1cda3cb30683f15b6a1a1dc4ea14a7" dependencies = [ "adler", ] [[package]] name = "mio" -version = "0.8.10" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f3d0b296e374a4e6f3c7b0a1f5a51d748a0d34c85e7dc48fc3fa9a87657fe09" +checksum = "a4a650543ca06a924e8b371db273b2756685faae30f8487da1b56505a8f78b0c" dependencies = [ "libc", "wasi", @@ -2260,6 +2751,12 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9252111cf132ba0929b6f8e030cac2a24b507f3a4d6db6fb2896f27b354c714b" +[[package]] +name = "murmurhash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2195bf6aa996a481483b29d62a7663eed3fe39600c460e323f8ff41e90bdd89b" + [[package]] name = "native-tls" version = "0.2.11" @@ -2288,6 +2785,30 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + +[[package]] +name = "num" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b05180d69e3da0e530ba2a1dae5110317e49e3b7f3d41be227dc5f92e49ee7af" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.4.4" @@ -2316,32 +2837,58 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-complex" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23c6602fda94a57c990fe0df199a035d83576b496aa29f4e634a8ac6004e68a6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + [[package]] name = "num-integer" -version = "0.1.45" +version = "0.1.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" dependencies = [ - "autocfg", "num-traits", ] [[package]] name = "num-iter" -version = "0.1.43" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d869c01cc0c455284163fd0092f1f93835385ccab5a98a0dcc497b2f8bf055a9" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252" +checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0" dependencies = [ "autocfg", + "num-bigint", "num-integer", "num-traits", ] [[package]] name = "num-traits" -version = "0.2.17" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" dependencies = [ "autocfg", "libm", @@ -2366,7 +2913,7 @@ dependencies = [ "base64 0.13.1", "chrono", "getrandom", - "http", + "http 0.2.12", "rand", "reqwest", "serde", @@ -2379,9 +2926,9 @@ dependencies = [ [[package]] name = "object" -version = "0.32.1" +version = "0.32.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9cf5f9dd3933bd50a9e1f149ec995f39ae2c496d31fd772c1fd45ebc27e902b0" +checksum = "a6a622008b6e321afc04970976f62ee297fdbaa6f95318ca343e3eebb9648441" dependencies = [ "memchr", ] @@ -2392,18 +2939,27 @@ version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +[[package]] +name = "oneshot" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f6640c6bda7731b1fdbab747981a0f896dd1fedaf9f4a53fa237a04a84431f4" +dependencies = [ + "loom", +] + [[package]] name = "openidconnect" -version = "3.4.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62d6050f6a84b81f23c569f5607ad883293e57491036e318fafe6fc4895fadb1" +checksum = "f47e80a9cfae4462dd29c41e987edd228971d6565553fbc14b8a11e666d91590" dependencies = [ "base64 0.13.1", "chrono", "dyn-clone", "ed25519-dalek", "hmac", - "http", + "http 0.2.12", "itertools 0.10.5", "log", "oauth2", @@ -2426,11 +2982,11 @@ dependencies = [ [[package]] name = "openssl" -version = "0.10.61" +version = "0.10.64" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b8419dc8cc6d866deb801274bba2e6f8f6108c1bb7fcc10ee5ab864931dbb45" +checksum = "95a0481286a310808298130d22dd1fef0fa571e05a8f44ec801801e84b216b1f" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.4.2", "cfg-if", "foreign-types", "libc", @@ -2447,7 +3003,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -2458,9 +3014,9 @@ checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" [[package]] name = "openssl-sys" -version = "0.9.97" +version = "0.9.101" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3eaad34cdd97d81de97964fc7f29e2d104f483840d906ef56daa1912338460b" +checksum = "dda2b0f344e78efc2facf7d195d098df0dd72151b26ab98da807afc26c198dff" dependencies = [ "cc", "libc", @@ -2483,6 +3039,21 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4030760ffd992bef45b0ae3f10ce1aba99e33464c90d14dd7c039884963ddc7a" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + +[[package]] +name = "ownedbytes" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e8a72b918ae8198abb3a18c190288123e1d442b6b9a7d709305fd194688b4b7" +dependencies = [ + "stable_deref_trait", +] + [[package]] name = "p256" version = "0.11.1" @@ -2549,9 +3120,9 @@ dependencies = [ [[package]] name = "pear" -version = "0.2.7" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61a386cd715229d399604b50d1361683fe687066f42d56f54be995bc6868f71c" +checksum = "4ccca0f6c17acc81df8e242ed473ec144cbf5c98037e69aa6d144780aad103c8" dependencies = [ "inlinable_string", "pear_codegen", @@ -2560,14 +3131,14 @@ dependencies = [ [[package]] name = "pear_codegen" -version = "0.2.7" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da9f0f13dac8069c139e8300a6510e3f4143ecf5259c60b116a9b271b4ca0d54" +checksum = "2e22670e8eb757cff11d6c199ca7b987f352f0346e0be4dd23869ec72cb53c77" dependencies = [ "proc-macro2", "proc-macro2-diagnostics", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -2576,7 +3147,7 @@ version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b8fcc794035347fb64beda2d3b462595dd2753e3f268d89c5aae77e8cf2c310" dependencies = [ - "base64 0.21.5", + "base64 0.21.7", "serde", ] @@ -2595,6 +3166,51 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "pest" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f8023d0fb78c8e03784ea1c7f3fa36e68a723138990b8d5a47d916b651e7a8" +dependencies = [ + "memchr", + "thiserror", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0d24f72393fd16ab6ac5738bc33cdb6a9aa73f8b902e8fe29cf4e67d7dd1026" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdc17e2a6c7d0a492f0158d7a4bd66cc17280308bbaff78d5bef566dca35ab80" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "pest_meta" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "934cd7631c050f4674352a6e835d5f6711ffbfb9345c2fc0107155ac495ae293" +dependencies = [ + "once_cell", + "pest", + "sha2", +] + [[package]] name = "petgraph" version = "0.6.4" @@ -2602,27 +3218,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1d3afd2628e69da2be385eb6f2fd57c8ac7977ceeff6dc166ff1657b0e386a9" dependencies = [ "fixedbitset", - "indexmap 2.1.0", + "indexmap 2.2.5", ] [[package]] name = "pin-project" -version = "1.1.3" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fda4ed1c6c173e3fc7a83629421152e01d7b1f9b7f65fb301e490e8cfc656422" +checksum = "b6bf43b791c5b9e34c3d182969b4abb522f9343702850a2e57f460d00d09b4b3" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.3" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4359fd9c9171ec6e8c62926d6faaf553a8dc3f64e1507e76da7911b4f6a04405" +checksum = "2f38a4412a78282e09a2cf38d195ea5420d15ba0602cb375210efbc877243965" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -2681,15 +3297,15 @@ dependencies = [ [[package]] name = "pkg-config" -version = "0.3.27" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" +checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec" [[package]] name = "platforms" -version = "3.2.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14e6ab3f592e6fb464fc9712d8d6e6912de6473954635fd76a589d832cffcbb0" +checksum = "626dec3cac7cc0e1577a2ec3fc496277ec2baa084bebad95bb6fdbfae235f84c" [[package]] name = "polling" @@ -2709,14 +3325,14 @@ dependencies = [ [[package]] name = "polling" -version = "3.3.1" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf63fa624ab313c11656b4cda960bfc46c410187ad493c41f6ba2d8c1e991c9e" +checksum = "24f040dee2588b4963afb4e420540439d126f73fdacf4a9c486a96d840bac3c9" dependencies = [ "cfg-if", "concurrent-queue", "pin-project-lite", - "rustix 0.38.28", + "rustix 0.38.31", "tracing", "windows-sys 0.52.0", ] @@ -2745,12 +3361,12 @@ dependencies = [ [[package]] name = "prettyplease" -version = "0.2.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae005bd773ab59b4725093fd7df83fd7892f7d8eafb48dbd7de6e024e4215f9d" +checksum = "a41cf62165e97c7f814d2221421dbb9afcbcdb0a88068e5ea206e19951c2cbb5" dependencies = [ "proc-macro2", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -2764,9 +3380,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.70" +version = "1.0.79" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39278fbbf5fb4f646ce651690877f89d1c5811a3d4acb27700c1cb3cdb78fd3b" +checksum = "e835ff2298f5721608eb1a980ecaee1aef2c132bf95ecc026a11b7bf3c01c02e" dependencies = [ "unicode-ident", ] @@ -2779,11 +3395,40 @@ checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", "version_check", "yansi", ] +[[package]] +name = "proptest" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b476131c3c86cb68032fdc5cb6d5a1045e3e42d96b69fa599fd77701e1f5bf" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags 2.4.2", + "lazy_static", + "num-traits", + "rand", + "rand_chacha", + "rand_xorshift", + "regex-syntax 0.8.2", + "rusty-fork", + "tempfile", + "unarray", +] + +[[package]] +name = "proptest-state-machine" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b52a714915de2d16a5289616d2265a934780f50a9dd30359322b687403fa2ac2" +dependencies = [ + "proptest", +] + [[package]] name = "prost" version = "0.11.9" @@ -2839,11 +3484,11 @@ dependencies = [ "multimap", "once_cell", "petgraph", - "prettyplease 0.2.15", + "prettyplease 0.2.16", "prost 0.12.3", "prost-types 0.12.3", "regex", - "syn 2.0.40", + "syn 2.0.52", "tempfile", "which", ] @@ -2871,7 +3516,7 @@ dependencies = [ "itertools 0.11.0", "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -2934,11 +3579,17 @@ dependencies = [ "zstd", ] +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quote" -version = "1.0.33" +version = "1.0.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae" +checksum = "291ec9ab5efd934aaf503a6466c5d5251535d108ee747472c3977cc5acc868ef" dependencies = [ "proc-macro2", ] @@ -2973,11 +3624,20 @@ dependencies = [ "getrandom", ] +[[package]] +name = "rand_xorshift" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f" +dependencies = [ + "rand_core", +] + [[package]] name = "rayon" -version = "1.8.0" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c27db03db7734835b3f53954b534c91069375ce6ccaa2e065441e07d9b6cdb1" +checksum = "e4963ed1bc86e4f3ee217022bd855b297cef07fb9eac5dfa1f788b220b49b3bd" dependencies = [ "either", "rayon-core", @@ -2985,9 +3645,9 @@ dependencies = [ [[package]] name = "rayon-core" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ce3fb6ad83f861aac485e76e1985cd109d9a3713802152be56c3b1f0e0658ed" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" dependencies = [ "crossbeam-deque", "crossbeam-utils", @@ -3004,25 +3664,34 @@ dependencies = [ [[package]] name = "regex" -version = "1.10.2" +version = "1.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "380b951a9c5e80ddfd6136919eef32310721aa4aacd4889a8d39124b026ab343" +checksum = "b62dbe01f0b06f9d8dc7d49e05a0785f153b00b2c227856282f671e0318c9b15" dependencies = [ "aho-corasick", "memchr", - "regex-automata", - "regex-syntax", + "regex-automata 0.4.6", + "regex-syntax 0.8.2", ] [[package]] name = "regex-automata" -version = "0.4.3" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +dependencies = [ + "regex-syntax 0.6.29", +] + +[[package]] +name = "regex-automata" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f804c7828047e88b2d32e2d7fe5a105da8ee3264f01902f796c8e067dc2483f" +checksum = "86b83b8b9847f9bf95ef68afb0b8e6cdb80f498442f5179a29fad448fcc1eaea" dependencies = [ "aho-corasick", "memchr", - "regex-syntax", + "regex-syntax 0.8.2", ] [[package]] @@ -3031,6 +3700,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "30b661b2f27137bdbc16f00eda72866a92bb28af1753ffbd56744fb6e2e9cd8e" +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + [[package]] name = "regex-syntax" version = "0.8.2" @@ -3039,17 +3714,17 @@ checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f" [[package]] name = "reqwest" -version = "0.11.22" +version = "0.11.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "046cd98826c46c2ac8ddecae268eb5c2e58628688a5fc7a2643704a73faba95b" +checksum = "78bf93c4af7a8bb7d879d51cebe797356ff10ae8516ace542b5182d9dcac10b2" dependencies = [ - "base64 0.21.5", + "base64 0.21.7", "bytes", "encoding_rs", "futures-core", "futures-util", "h2", - "http", + "http 0.2.12", "http-body", "hyper", "hyper-rustls", @@ -3065,6 +3740,7 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", + "sync_wrapper", "system-configuration", "tokio", "tokio-rustls", @@ -3100,16 +3776,27 @@ dependencies = [ [[package]] name = "ring" -version = "0.17.7" +version = "0.17.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "688c63d65483050968b2a8937f7995f443e27041a0f7700aa59b0822aedebb74" +checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d" dependencies = [ "cc", + "cfg-if", "getrandom", "libc", "spin 0.9.8", "untrusted", - "windows-sys 0.48.0", + "windows-sys 0.52.0", +] + +[[package]] +name = "roaring" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1c77081a55300e016cb86f2864415b7518741879db925b8d488a0ee0d2da6bf" +dependencies = [ + "bytemuck", + "byteorder", ] [[package]] @@ -3132,12 +3819,28 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust-stemmers" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e46a2036019fdb888131db7a4c847a1063a7493f971ed94ea82c67eada63ca54" +dependencies = [ + "serde", + "serde_derive", +] + [[package]] name = "rustc-demangle" version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc_version" version = "0.4.0" @@ -3163,14 +3866,14 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.28" +version = "0.38.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72e572a5e8ca657d7366229cdde4bd14c4eb5499a9573d4d366fe1b599daa316" +checksum = "6ea3e1a662af26cd7a3ba09c0297a31af215563ecf42817c98df621387f4e949" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.4.2", "errno", "libc", - "linux-raw-sys 0.4.12", + "linux-raw-sys 0.4.13", "windows-sys 0.52.0", ] @@ -3204,7 +3907,7 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c74cae0a4cf6ccbbf5f359f08efdf8ee7e1dc532573bf0db71968cb56b1448c" dependencies = [ - "base64 0.21.5", + "base64 0.21.7", ] [[package]] @@ -3223,19 +3926,31 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ffc183a10b4478d04cbbbfc96d0873219d962dd5accaff2ffbd4ceb7df837f4" +[[package]] +name = "rusty-fork" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3dcc6e454c328bb824492db107ab7c0ae8fcffe4ad210136ef014458c1bc4f" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + [[package]] name = "ryu" -version = "1.0.16" +version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f98d2aa92eebf49b69786be48e4477826b256916e84a57ff2a4f21923b48eb4c" +checksum = "e86697c916019a8588c99b5fac3cead74ec0b4b819707a682fd4d23fa0ce1ba1" [[package]] name = "schannel" -version = "0.1.22" +version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c3733bf4cf7ea0880754e19cb5a462007c4a8c1914bff372ccc95b464f1df88" +checksum = "fbc91545643bcf3a0bbb6569265615222618bdf33ce4ffbbd13c4bbd4c093534" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.52.0", ] [[package]] @@ -3262,6 +3977,12 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "scoped-tls" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" + [[package]] name = "scopeguard" version = "1.2.0" @@ -3341,15 +4062,15 @@ dependencies = [ [[package]] name = "semver" -version = "1.0.20" +version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836fa6a3e1e547f9a2c4040802ec865b5d85f4014efe00555d7090a3dcaa1090" +checksum = "92d43fe69e652f3df9bdc2b85b2854a0825b86e4fb76bc44d945137d053639ca" [[package]] name = "serde" -version = "1.0.193" +version = "1.0.197" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25dd9975e68d0cb5aa1120c288333fc98731bd1dd12f561e468ea4728c042b89" +checksum = "3fb1c873e1b9b056a4dc4c0c198b24c3ffa059243875552b2bd0933b1aee4ce2" dependencies = [ "serde_derive", ] @@ -3366,13 +4087,13 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.193" +version = "1.0.197" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43576ca501357b9b071ac53cdc7da8ef0cbd9493d8df094cd821777ea6e894d3" +checksum = "7eb0b34b42edc17f6b7cac84a52a1c5f0e1bb2227e997ca9011ea3dd34e8610b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -3388,11 +4109,10 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.108" +version = "1.0.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d1c7e3eac408d115102c4c24ad393e0821bb3a5df4d506a80f85f7a742a526b" +checksum = "c5f09b1bd632ef549eaa9f60a1f8de742bdbc698e6cee2095fc84dde5f549ae0" dependencies = [ - "indexmap 2.1.0", "itoa", "ryu", "serde", @@ -3400,9 +4120,9 @@ dependencies = [ [[package]] name = "serde_path_to_error" -version = "0.1.14" +version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4beec8bce849d58d06238cb50db2e1c417cfeafa4c63f692b15c82b7c80f8335" +checksum = "af99884400da37c88f5e9146b7f1fd0fbcae8f6eec4e9da38b67d05486f814a6" dependencies = [ "itoa", "serde", @@ -3431,16 +4151,17 @@ dependencies = [ [[package]] name = "serde_with" -version = "3.4.0" +version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64cd236ccc1b7a29e7e2739f27c0b2dd199804abc4290e32f59f3b68d6405c23" +checksum = "ee80b0e361bbf88fd2f6e242ccd19cfda072cb0faa6ae694ecee08199938569a" dependencies = [ - "base64 0.21.5", + "base64 0.21.7", "chrono", "hex", "indexmap 1.9.3", - "indexmap 2.1.0", + "indexmap 2.2.5", "serde", + "serde_derive", "serde_json", "serde_with_macros", "time", @@ -3448,23 +4169,23 @@ dependencies = [ [[package]] name = "serde_with_macros" -version = "3.4.0" +version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93634eb5f75a2323b16de4748022ac4297f9e76b6dced2be287a099f41b5e788" +checksum = "6561dc161a9224638a31d876ccdfefbc1df91d3f3a8342eddb35f055d48c7655" dependencies = [ "darling", "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] name = "serde_yaml" -version = "0.9.27" +version = "0.9.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cc7a1570e38322cfe4154732e5110f887ea57e22b76f4bfd32b5bdd3368666c" +checksum = "8fd075d994154d4a774f95b51fb96bdc2832b0ea48425c92546073816cda1f2f" dependencies = [ - "indexmap 2.1.0", + "indexmap 2.2.5", "itoa", "ryu", "serde", @@ -3493,6 +4214,15 @@ dependencies = [ "digest", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "signal-hook-registry" version = "1.4.1" @@ -3522,6 +4252,15 @@ dependencies = [ "rand_core", ] +[[package]] +name = "sketches-ddsketch" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85636c14b73d81f541e525f585c0a2109e6744e1565b5c1668e31c70c10ed65c" +dependencies = [ + "serde", +] + [[package]] name = "slab" version = "0.4.9" @@ -3533,9 +4272,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.11.2" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dccd0940a2dcdf68d092b8cbab7dc0ad8fa938bf95787e1b916b0e3d0e8e970" +checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7" [[package]] name = "snap" @@ -3555,12 +4294,12 @@ dependencies = [ [[package]] name = "socket2" -version = "0.5.5" +version = "0.5.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b5fac59a5cb5dd637972e5fca70daf0523c9067fcdc4842f053dae04a18f8e9" +checksum = "05ffd9c0a93b7543e062e759284fcf5f5e3b098501104bfbdde4d404db792871" dependencies = [ "libc", - "windows-sys 0.48.0", + "windows-sys 0.52.0", ] [[package]] @@ -3595,6 +4334,18 @@ dependencies = [ "der 0.7.8", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "strsim" version = "0.10.0" @@ -3620,9 +4371,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.40" +version = "2.0.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13fa70a4ee923979ffb522cacce59d34421ebdea5625e1073c4326ef9d2dd42e" +checksum = "b699d15b36d1f02c3e7c69f8ffef53de37aefae075d8488d4ba1a7788d574a07" dependencies = [ "proc-macro2", "quote", @@ -3656,47 +4407,197 @@ dependencies = [ "libc", ] +[[package]] +name = "tantivy" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6083cd777fa94271b8ce0fe4533772cb8110c3044bab048d20f70108329a1f2" +dependencies = [ + "aho-corasick", + "arc-swap", + "async-trait", + "base64 0.21.7", + "bitpacking", + "byteorder", + "census", + "crc32fast", + "crossbeam-channel", + "downcast-rs", + "fastdivide", + "fs4", + "htmlescape", + "itertools 0.11.0", + "levenshtein_automata", + "log", + "lru", + "lz4_flex", + "measure_time", + "memmap2", + "murmurhash32", + "num_cpus", + "once_cell", + "oneshot", + "rayon", + "regex", + "rust-stemmers", + "rustc-hash", + "serde", + "serde_json", + "sketches-ddsketch", + "smallvec", + "tantivy-bitpacker", + "tantivy-columnar", + "tantivy-common", + "tantivy-fst", + "tantivy-query-grammar", + "tantivy-stacker", + "tantivy-tokenizer-api", + "tempfile", + "thiserror", + "time", + "uuid", + "winapi", +] + +[[package]] +name = "tantivy-bitpacker" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecb164321482301f514dd582264fa67f70da2d7eb01872ccd71e35e0d96655a" +dependencies = [ + "bitpacking", +] + +[[package]] +name = "tantivy-columnar" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d85f8019af9a78b3118c11298b36ffd21c2314bd76bbcd9d12e00124cbb7e70" +dependencies = [ + "fastdivide", + "fnv", + "itertools 0.11.0", + "serde", + "tantivy-bitpacker", + "tantivy-common", + "tantivy-sstable", + "tantivy-stacker", +] + +[[package]] +name = "tantivy-common" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af4a3a975e604a2aba6b1106a04505e1e7a025e6def477fab6e410b4126471e1" +dependencies = [ + "async-trait", + "byteorder", + "ownedbytes", + "serde", + "time", +] + +[[package]] +name = "tantivy-fst" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc3c506b1a8443a3a65352df6382a1fb6a7afe1a02e871cee0d25e2c3d5f3944" +dependencies = [ + "byteorder", + "regex-syntax 0.6.29", + "utf8-ranges", +] + +[[package]] +name = "tantivy-query-grammar" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d39c5a03100ac10c96e0c8b07538e2ab8b17da56434ab348309b31f23fada77" +dependencies = [ + "nom", +] + +[[package]] +name = "tantivy-sstable" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0c1bb43e5e8b8e05eb8009610344dbf285f06066c844032fbb3e546b3c71df" +dependencies = [ + "tantivy-common", + "tantivy-fst", + "zstd", +] + +[[package]] +name = "tantivy-stacker" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2c078595413f13f218cf6f97b23dcfd48936838f1d3d13a1016e05acd64ed6c" +dependencies = [ + "murmurhash32", + "tantivy-common", +] + +[[package]] +name = "tantivy-tokenizer-api" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "347b6fb212b26d3505d224f438e3c4b827ab8bd847fe9953ad5ac6b8f9443b66" +dependencies = [ + "serde", +] + [[package]] name = "tempfile" -version = "3.8.1" +version = "3.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ef1adac450ad7f4b3c28589471ade84f25f731a7a0fe30d71dfa9f60fd808e5" +checksum = "85b77fafb263dd9d05cbeac119526425676db3784113aa9295c88498cbf8bff1" dependencies = [ "cfg-if", "fastrand 2.0.1", - "redox_syscall", - "rustix 0.38.28", - "windows-sys 0.48.0", + "rustix 0.38.31", + "windows-sys 0.52.0", ] [[package]] name = "thiserror" -version = "1.0.50" +version = "1.0.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9a7210f5c9a7156bb50aa36aed4c95afb51df0df00713949448cf9e97d382d2" +checksum = "03468839009160513471e86a034bb2c5c0e4baae3b43f79ffc55c4a5427b3297" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.50" +version = "1.0.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "266b2e40bc00e5a6c09c3584011e08b06f123c00362c92b975ba9843aaaa14b8" +checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", +] + +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", ] [[package]] name = "time" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4a34ab300f2dee6e562c10a046fc05e358b29f9bf92277f30c3c8d82275f6f5" +checksum = "c8248b6521bb14bc45b4067159b9b6ad792e2d6d754d6c41fb50e29fefe38749" dependencies = [ "deranged", "itoa", + "num-conv", "powerfmt", "serde", "time-core", @@ -3711,13 +4612,23 @@ checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3" [[package]] name = "time-macros" -version = "0.2.15" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ad70d68dba9e1f8aceda7aa6711965dfec1cac869f311a51bd08b3a2ccbce20" +checksum = "7ba3a3ef41e6672a2f0f001392bb5dcd3ff0a9992d618ca761a11c3121547774" dependencies = [ + "num-conv", "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinyvec" version = "1.6.0" @@ -3735,9 +4646,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.35.0" +version = "1.36.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "841d45b238a16291a4e1584e61820b8ae57d696cc5015c459c229ccc6990cc1c" +checksum = "61285f6515fa018fb2d1e46eb21223fff441ee8db5d0f1435e8ab4f5cdb80931" dependencies = [ "backtrace", "bytes", @@ -3746,7 +4657,7 @@ dependencies = [ "num_cpus", "pin-project-lite", "signal-hook-registry", - "socket2 0.5.5", + "socket2 0.5.6", "tokio-macros", "windows-sys 0.48.0", ] @@ -3769,7 +4680,7 @@ checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -3827,10 +4738,10 @@ dependencies = [ "async-stream", "async-trait", "axum", - "base64 0.21.5", + "base64 0.21.7", "bytes", "h2", - "http", + "http 0.2.12", "http-body", "hyper", "hyper-timeout", @@ -3851,11 +4762,11 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d021fc044c18582b9a2408cd0dd05b1596e3ecdb5c4df822bb0183545683889" dependencies = [ - "prettyplease 0.2.15", + "prettyplease 0.2.16", "proc-macro2", "prost-build 0.12.3", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -3884,12 +4795,12 @@ version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c5bb1d698276a2443e5ecfabc1008bf15a36c12e6a7176e7bf089ea9131140" dependencies = [ - "base64 0.21.5", - "bitflags 2.4.1", + "base64 0.21.7", + "bitflags 2.4.2", "bytes", "futures-core", "futures-util", - "http", + "http 0.2.12", "http-body", "http-range-header", "mime", @@ -3931,7 +4842,7 @@ checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] @@ -3941,13 +4852,43 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", ] [[package]] name = "treediff" -version = "4.0.2" +version = "4.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52984d277bdf2a751072b5df30ec0377febdb02f7696d64c2d7d54630bac4303" +checksum = "4d127780145176e2b5d16611cc25a900150e86e9fd79d3bde6ff3a37359c9cb5" dependencies = [ "serde_json", ] @@ -3964,20 +4905,32 @@ version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +[[package]] +name = "ucd-trie" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9" + +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "uncased" -version = "0.9.9" +version = "0.9.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b9bc53168a4be7402ab86c3aad243a84dd7381d09be0eddc81280c1da95ca68" +checksum = "e1b88fcfe09e89d3866a5c11019378088af2d24c3fbd4f0543f96b479ec90697" dependencies = [ "version_check", ] [[package]] name = "unicode-bidi" -version = "0.3.14" +version = "0.3.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f2528f27a9eb2b21e69c95319b30bd0efd85d09c379741b0f78ea1d86be2416" +checksum = "08f95100a766bf4f8f28f90d77e0a5461bbdb219042e7679bebe79004fed8d75" [[package]] name = "unicode-ident" @@ -3987,18 +4940,18 @@ checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" [[package]] name = "unicode-normalization" -version = "0.1.22" +version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c5713f0fc4b5db668a2ac63cdb7bb4469d8c9fed047b1d0292cc7b0ce2ba921" +checksum = "a56d1686db2308d901306f92a263857ef59ea39678a5458e7cb17f01415101f5" dependencies = [ "tinyvec", ] [[package]] name = "unsafe-libyaml" -version = "0.2.9" +version = "0.2.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f28467d3e1d3c6586d8f25fa243f544f5800fec42d97032474e17222c2b75cfa" +checksum = "ab4c90930b95a82d00dc9e9ac071b4991924390d46cbd0dfe566148667605e4b" [[package]] name = "untrusted" @@ -4024,33 +4977,46 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" +[[package]] +name = "utf8-ranges" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fcfc827f90e53a02eaef5e535ee14266c1d569214c6aa70133a624d8a3164ba" + [[package]] name = "uuid" -version = "1.6.1" +version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e395fcf16a7a3d8127ec99782007af141946b4795001f876d54fb0d55978560" +checksum = "f00cc9702ca12d3c81455259621e676d0f7251cec66a21e98fe2e9a37db93b2a" dependencies = [ "getrandom", "rand", + "serde", "uuid-macro-internal", ] [[package]] name = "uuid-macro-internal" -version = "1.6.1" +version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f49e7f3f3db8040a100710a11932239fd30697115e2ba4107080d8252939845e" +checksum = "7abb14ae1a50dad63eaa768a458ef43d298cd1bd44951677bd10b732a9ba2a2d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "value-bag" -version = "1.4.2" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a72e1902dde2bd6441347de2b70b7f5d59bf157c6c62f0c44572607a1d55bbe" +checksum = "8fec26a25bd6fca441cdd0f769fd7f891bae119f996de31f86a5eddccef54c1d" [[package]] name = "vcpkg" @@ -4070,6 +5036,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "wait-timeout" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f200f5b12eb75f8c1ed65abd4b2db8a6e1b138a20de009dacee265a2498f3f6" +dependencies = [ + "libc", +] + [[package]] name = "waker-fn" version = "1.1.1" @@ -4093,9 +5068,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wasm-bindgen" -version = "0.2.89" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ed0d4f68a3015cc185aff4db9506a015f4b96f95303897bfa23f846db54064e" +checksum = "4be2531df63900aeb2bca0daaaddec08491ee64ceecbee5076636a3b026795a8" dependencies = [ "cfg-if", "wasm-bindgen-macro", @@ -4103,24 +5078,24 @@ dependencies = [ [[package]] name = "wasm-bindgen-backend" -version = "0.2.89" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b56f625e64f3a1084ded111c4d5f477df9f8c92df113852fa5a374dbda78826" +checksum = "614d787b966d3989fa7bb98a654e369c762374fd3213d212cfc0251257e747da" dependencies = [ "bumpalo", "log", "once_cell", "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-futures" -version = "0.4.39" +version = "0.4.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac36a15a220124ac510204aec1c3e5db8a22ab06fd6706d881dc6149f8ed9a12" +checksum = "76bc14366121efc8dbb487ab05bcc9d346b3b5ec0eaa76e46594cabbe51762c0" dependencies = [ "cfg-if", "js-sys", @@ -4130,9 +5105,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.89" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0162dbf37223cd2afce98f3d0785506dcb8d266223983e4b5b525859e6e182b2" +checksum = "a1f8823de937b71b9460c0c34e25f3da88250760bec0ebac694b49997550d726" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4140,28 +5115,28 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.89" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0eb82fcb7930ae6219a7ecfd55b217f5f0893484b7a13022ebb2b2bf20b5283" +checksum = "e94f17b526d0a461a191c78ea52bbce64071ed5c04c9ffe424dcb38f74171bb7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", "wasm-bindgen-backend", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.89" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ab9b36309365056cd639da3134bf87fa8f3d86008abf99e612384a6eecd459f" +checksum = "af190c94f2773fdb3729c55b007a722abb5384da03bc0986df4c289bf5567e96" [[package]] name = "web-sys" -version = "0.3.66" +version = "0.3.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50c24a44ec86bb68fbecd1b3efed7e85ea5621b39b35ef2766b66cd984f8010f" +checksum = "77afa9a11836342370f4817622a2f0f418b134426d91a82dfb48f532d2ec13ef" dependencies = [ "js-sys", "wasm-bindgen", @@ -4169,9 +5144,9 @@ dependencies = [ [[package]] name = "webpki-roots" -version = "0.25.3" +version = "0.25.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1778a42e8b3b90bff8d0f5032bf22250792889a5cdc752aa0020c84abe3aaf10" +checksum = "5f20c57d8d7db6d3b86154206ae5d8fba62dd39573114de97c2cb0578251f8e1" [[package]] name = "which" @@ -4182,7 +5157,7 @@ dependencies = [ "either", "home", "once_cell", - "rustix 0.38.28", + "rustix 0.38.31", ] [[package]] @@ -4208,14 +5183,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] -name = "windows-core" -version = "0.51.1" +name = "windows" +version = "0.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1f8cf84f35d2db49a46868f947758c7a1138116f7fac3bc844f43ade1292e64" +checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f" dependencies = [ "windows-targets 0.48.5", ] +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets 0.52.4", +] + [[package]] name = "windows-sys" version = "0.48.0" @@ -4231,7 +5215,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.0", + "windows-targets 0.52.4", ] [[package]] @@ -4251,17 +5235,17 @@ dependencies = [ [[package]] name = "windows-targets" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a18201040b24831fbb9e4eb208f8892e1f50a37feb53cc7ff887feb8f50e7cd" +checksum = "7dd37b7e5ab9018759f893a1952c9420d060016fc19a472b4bb20d1bdd694d1b" dependencies = [ - "windows_aarch64_gnullvm 0.52.0", - "windows_aarch64_msvc 0.52.0", - "windows_i686_gnu 0.52.0", - "windows_i686_msvc 0.52.0", - "windows_x86_64_gnu 0.52.0", - "windows_x86_64_gnullvm 0.52.0", - "windows_x86_64_msvc 0.52.0", + "windows_aarch64_gnullvm 0.52.4", + "windows_aarch64_msvc 0.52.4", + "windows_i686_gnu 0.52.4", + "windows_i686_msvc 0.52.4", + "windows_x86_64_gnu 0.52.4", + "windows_x86_64_gnullvm 0.52.4", + "windows_x86_64_msvc 0.52.4", ] [[package]] @@ -4272,9 +5256,9 @@ checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" [[package]] name = "windows_aarch64_gnullvm" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb7764e35d4db8a7921e09562a0304bf2f93e0a51bfccee0bd0bb0b666b015ea" +checksum = "bcf46cf4c365c6f2d1cc93ce535f2c8b244591df96ceee75d8e83deb70a9cac9" [[package]] name = "windows_aarch64_msvc" @@ -4284,9 +5268,9 @@ checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" [[package]] name = "windows_aarch64_msvc" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbaa0368d4f1d2aaefc55b6fcfee13f41544ddf36801e793edbbfd7d7df075ef" +checksum = "da9f259dd3bcf6990b55bffd094c4f7235817ba4ceebde8e6d11cd0c5633b675" [[package]] name = "windows_i686_gnu" @@ -4296,9 +5280,9 @@ checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" [[package]] name = "windows_i686_gnu" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a28637cb1fa3560a16915793afb20081aba2c92ee8af57b4d5f28e4b3e7df313" +checksum = "b474d8268f99e0995f25b9f095bc7434632601028cf86590aea5c8a5cb7801d3" [[package]] name = "windows_i686_msvc" @@ -4308,9 +5292,9 @@ checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" [[package]] name = "windows_i686_msvc" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffe5e8e31046ce6230cc7215707b816e339ff4d4d67c65dffa206fd0f7aa7b9a" +checksum = "1515e9a29e5bed743cb4415a9ecf5dfca648ce85ee42e15873c3cd8610ff8e02" [[package]] name = "windows_x86_64_gnu" @@ -4320,9 +5304,9 @@ checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" [[package]] name = "windows_x86_64_gnu" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d6fa32db2bc4a2f5abeacf2b69f7992cd09dca97498da74a151a3132c26befd" +checksum = "5eee091590e89cc02ad514ffe3ead9eb6b660aedca2183455434b93546371a03" [[package]] name = "windows_x86_64_gnullvm" @@ -4332,9 +5316,9 @@ checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" [[package]] name = "windows_x86_64_gnullvm" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a657e1e9d3f514745a572a6846d3c7aa7dbe1658c056ed9c3344c4109a6949e" +checksum = "77ca79f2451b49fa9e2af39f0747fe999fcda4f5e241b2898624dca97a1f2177" [[package]] name = "windows_x86_64_msvc" @@ -4344,9 +5328,9 @@ checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" [[package]] name = "windows_x86_64_msvc" -version = "0.52.0" +version = "0.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dff9641d1cd4be8d1a070daf9e3773c5f67e78b4d9d42263020c057706765c04" +checksum = "32b752e52a2da0ddfbdbcc6fceadfeede4c939ed16d13e648833a61dfb611ed8" [[package]] name = "winreg" @@ -4362,6 +5346,7 @@ dependencies = [ name = "worker" version = "0.1.0" dependencies = [ + "arrow", "async-trait", "aws-config", "aws-sdk-s3", @@ -4373,17 +5358,20 @@ dependencies = [ "k8s-openapi", "kube", "murmur3", - "num-bigint", "num_cpus", "parking_lot", + "proptest", + "proptest-state-machine", "prost 0.12.3", "prost-types 0.12.3", "pulsar", "rand", "rayon", + "roaring", "schemars", "serde", "serde_json", + "tantivy", "tempfile", "thiserror", "tokio", @@ -4401,28 +5389,28 @@ checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" [[package]] name = "yansi" -version = "1.0.0-rc.1" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1367295b8f788d371ce2dbc842c7b709c73ee1364d30351dd300ec2203b12377" +checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" [[package]] name = "zerocopy" -version = "0.7.30" +version = "0.7.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "306dca4455518f1f31635ec308b6b3e4eb1b11758cefafc782827d0aa7acb5c7" +checksum = "74d4d3961e53fa4c9a25a8637fc2bfaf2595b3d3ae34875568a5cf64787716be" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.7.30" +version = "0.7.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be912bf68235a88fbefd1b73415cb218405958d1655b2ece9035a19920bdf6ba" +checksum = "9ce1b18ccd8e73a9321186f97e46f9f04b778851177567b1975109d26a08d2a6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.40", + "syn 2.0.52", ] [[package]] diff --git a/DEVELOP.md b/DEVELOP.md index f034e07bed3..c9550e639f4 100644 --- a/DEVELOP.md +++ b/DEVELOP.md @@ -6,8 +6,6 @@ https://packaging.python.org. ## Setup -Because of the dependencies it relies on (like `pytorch`), this project does not support Python version >3.10.0. - Set up a virtual environment and install the project's requirements and dev requirements: @@ -50,6 +48,15 @@ api = chromadb.HttpClient(host="localhost", port="8000") print(api.heartbeat()) ``` +## Local dev setup for distributed chroma +We use tilt for providing local dev setup. Tilt is an open source project +##### Requirement +- Docker +- Local Kubernetes cluster (Recommended: [OrbStack](https://orbstack.dev/) for mac, [Kind](https://kind.sigs.k8s.io/) for linux) +- [Tilt](https://docs.tilt.dev/) + +For starting the distributed Chroma in the workspace, use `tilt up`. It will create all the required resources and build the necessary Docker image in the current kubectl context. +Once done, it will expose Chroma on port 8000. You can also visit the Tilt dashboard UI at http://localhost:10350/. To clean and remove all the resources created by Tilt, use `tilt down`. ## Testing diff --git a/Dockerfile b/Dockerfile index 1f90733edbb..d12d5620fc1 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,6 +25,17 @@ COPY --from=builder /install /usr/local COPY ./bin/docker_entrypoint.sh /docker_entrypoint.sh COPY ./ /chroma +RUN apt-get update --fix-missing && apt-get install -y curl && \ + chmod +x /docker_entrypoint.sh && \ + rm -rf /var/lib/apt/lists/* + +ENV CHROMA_HOST_ADDR "0.0.0.0" +ENV CHROMA_HOST_PORT 8000 +ENV CHROMA_WORKERS 1 +ENV CHROMA_LOG_CONFIG "chromadb/log_config.yml" +ENV CHROMA_TIMEOUT_KEEP_ALIVE 30 + EXPOSE 8000 -CMD ["/docker_entrypoint.sh"] +ENTRYPOINT ["/docker_entrypoint.sh"] +CMD [ "--workers ${CHROMA_WORKERS} --host ${CHROMA_HOST_ADDR} --port ${CHROMA_HOST_PORT} --proxy-headers --log-config ${CHROMA_LOG_CONFIG} --timeout-keep-alive ${CHROMA_TIMEOUT_KEEP_ALIVE}"] \ No newline at end of file diff --git a/Tiltfile b/Tiltfile new file mode 100644 index 00000000000..fce67deacc6 --- /dev/null +++ b/Tiltfile @@ -0,0 +1,121 @@ +update_settings(max_parallel_updates=6) + +docker_build( + 'local:sysdb-migration', + '.', + only=['go/'], + dockerfile='./go/Dockerfile.migration' +) + +docker_build( + 'local:sysdb', + '.', + only=['go/', 'idl/'], + dockerfile='./go/Dockerfile' +) + +docker_build( + 'local:frontend-service', + '.', + only=['chromadb/', 'idl/', 'requirements.txt', 'bin/'], + dockerfile='./Dockerfile', +) + +docker_build( + 'local:query-service', + '.', + only=["rust/", "idl/", "Cargo.toml", "Cargo.lock"], + dockerfile='./rust/worker/Dockerfile', + target='query_service' +) + +docker_build( + 'local:compaction-service', + '.', + only=["rust/", "idl/", "Cargo.toml", "Cargo.lock"], + dockerfile='./rust/worker/Dockerfile', + target='compaction_service' +) + +k8s_yaml( + helm( + 'k8s/distributed-chroma', + namespace='chroma', + values=[ + 'k8s/distributed-chroma/values.yaml' + ] + ) +) + +k8s_yaml([ + 'k8s/test/postgres.yaml', +]) + +# Extra stuff to make debugging and testing easier +k8s_yaml([ + 'k8s/test/sysdb-service.yaml', + 'k8s/test/jaeger-service.yaml', + 'k8s/test/pulsar-service.yaml', + 'k8s/test/logservice-service.yaml', + 'k8s/test/minio.yaml', + 'k8s/test/query-service-service.yaml', + 'k8s/test/test-memberlist-cr.yaml', +]) + +# Lots of things assume the cluster is in a basic state. Get it into a basic +# state before deploying anything else. +k8s_resource( + objects=['chroma:Namespace'], + new_name='namespace', + labels=["infrastructure"], +) +k8s_resource( + objects=[ + 'pod-watcher:Role', + 'memberlists.chroma.cluster:CustomResourceDefinition', + 'query-service-memberlist:MemberList', + 'compaction-service-memberlist:MemberList', + + 'sysdb-serviceaccount:serviceaccount', + 'sysdb-serviceaccount-rolebinding:RoleBinding', + 'sysdb-query-service-memberlist-binding:clusterrolebinding', + + 'logservice-serviceaccount:serviceaccount', + + 'query-service-serviceaccount:serviceaccount', + 'query-service-serviceaccount-rolebinding:RoleBinding', + 'query-service-memberlist-readerwriter:ClusterRole', + 'query-service-query-service-memberlist-binding:clusterrolebinding', + 'query-service-memberlist-readerwriter-binding:clusterrolebinding', + + 'compaction-service-memberlist-readerwriter:ClusterRole', + 'compaction-service-compaction-service-memberlist-binding:clusterrolebinding', + 'compaction-service-memberlist-readerwriter-binding:clusterrolebinding', + 'compaction-service-serviceaccount:serviceaccount', + 'compaction-service-serviceaccount-rolebinding:RoleBinding', + + 'test-memberlist:MemberList', + 'test-memberlist-reader:ClusterRole', + 'test-memberlist-reader-binding:ClusterRoleBinding', + ], + new_name='k8s_setup', + labels=["infrastructure"], + resource_deps=['namespace'], +) + +# Production Chroma +k8s_resource('postgres', resource_deps=['k8s_setup', 'namespace'], labels=["infrastructure"]) +k8s_resource('pulsar', resource_deps=['k8s_setup', 'namespace'], labels=["infrastructure"], port_forwards=['6650:6650', '8080:8080']) +k8s_resource('sysdb-migration', resource_deps=['postgres', 'namespace'], labels=["infrastructure"]) +k8s_resource('logservice', resource_deps=['sysdb-migration'], labels=["chroma"], port_forwards='50052:50051') +k8s_resource('sysdb', resource_deps=['pulsar', 'sysdb-migration'], labels=["chroma"], port_forwards='50051:50051') +k8s_resource('frontend-service', resource_deps=['pulsar', 'sysdb', 'logservice'],labels=["chroma"], port_forwards='8000:8000') +k8s_resource('query-service', resource_deps=['sysdb'], labels=["chroma"]) +k8s_resource('compaction-service', resource_deps=['sysdb'], labels=["chroma"]) + +# I have no idea why these need their own lines but the others don't. +k8s_resource(objects=['query-service:service'], new_name='query-service-service', resource_deps=['query-service'], labels=["chroma"]) +k8s_resource(objects=['jaeger-lb:Service'], new_name='jaeger-service', resource_deps=['k8s_setup'], labels=["debug"]) + +# Local S3 +k8s_resource('minio-deployment', resource_deps=['k8s_setup'], labels=["debug"], port_forwards='9000:9000') diff --git a/bin/cluster-test.sh b/bin/cluster-test.sh index 10c48781c07..5f86a78e2c6 100755 --- a/bin/cluster-test.sh +++ b/bin/cluster-test.sh @@ -1,62 +1,20 @@ #!/usr/bin/env bash - set -e -function cleanup { - # Restore the previous kube context - kubectl config use-context $PREV_CHROMA_KUBE_CONTEXT - # Kill the tunnel process - kill $TUNNEL_PID - minikube delete -p chroma-test -} - -trap cleanup EXIT - -# Save the current kube context into a variable -export PREV_CHROMA_KUBE_CONTEXT=$(kubectl config current-context) - -# Create a new minikube cluster for the test -minikube start -p chroma-test - -# Add the ingress addon to the cluster -minikube addons enable ingress -p chroma-test -minikube addons enable ingress-dns -p chroma-test - -# Setup docker to build inside the minikube cluster and build the image -eval $(minikube -p chroma-test docker-env) -docker build -t server:latest -f Dockerfile . -docker build -t chroma-coordinator:latest -f go/coordinator/Dockerfile . -docker build -t worker -f rust/worker/Dockerfile . --build-arg CHROMA_KUBERNETES_INTEGRATION=1 - -# Apply the kubernetes manifests -kubectl apply -f k8s/deployment -kubectl apply -f k8s/crd -kubectl apply -f k8s/cr -kubectl apply -f k8s/test - -# Wait for the pods in the chroma namespace to be ready -kubectl wait --namespace chroma --for=condition=Ready pods --all --timeout=400s - -# Run mini kube tunnel in the background to expose the service -minikube tunnel -c true -p chroma-test & -TUNNEL_PID=$! - -# Wait for the tunnel to be ready. There isn't an easy way to check if the tunnel is ready. So we just wait for 10 seconds -sleep 10 - +# TODO make url configuration consistent. export CHROMA_CLUSTER_TEST_ONLY=1 -export CHROMA_SERVER_HOST=$(kubectl get svc server -n chroma -o=jsonpath='{.status.loadBalancer.ingress[0].ip}') -export PULSAR_BROKER_URL=$(kubectl get svc pulsar-lb -n chroma -o=jsonpath='{.status.loadBalancer.ingress[0].ip}') -export CHROMA_COORDINATOR_HOST=$(kubectl get svc coordinator-lb -n chroma -o=jsonpath='{.status.loadBalancer.ingress[0].ip}') -export CHROMA_SERVER_GRPC_PORT="50051" +export CHROMA_SERVER_HOST=localhost:8000 +export PULSAR_BROKER_URL=localhost +export CHROMA_COORDINATOR_HOST=localhost echo "Chroma Server is running at port $CHROMA_SERVER_HOST" echo "Pulsar Broker is running at port $PULSAR_BROKER_URL" echo "Chroma Coordinator is running at port $CHROMA_COORDINATOR_HOST" -echo testing: python -m pytest "$@" -python -m pytest "$@" +kubectl -n chroma port-forward svc/sysdb-lb 50051:50051 & +kubectl -n chroma port-forward svc/logservice-lb 50052:50051 & +kubectl -n chroma port-forward svc/pulsar-lb 6650:6650 & +kubectl -n chroma port-forward svc/pulsar-lb 8080:8080 & +kubectl -n chroma port-forward svc/frontend-service 8000:8000 & -export CHROMA_KUBERNETES_INTEGRATION=1 -cd go/coordinator -go test -timeout 30s -run ^TestNodeWatcher$ github.com/chroma/chroma-coordinator/internal/memberlist_manager +"$@" diff --git a/bin/docker_entrypoint.sh b/bin/docker_entrypoint.sh index e6f2df70be8..e9498b4fd7c 100755 --- a/bin/docker_entrypoint.sh +++ b/bin/docker_entrypoint.sh @@ -1,5 +1,15 @@ #!/bin/bash +set -e export IS_PERSISTENT=1 export CHROMA_SERVER_NOFILE=65535 -exec uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30 +args="$@" + +if [[ $args =~ ^uvicorn.* ]]; then + echo "Starting server with args: $(eval echo "$args")" + echo -e "\033[31mWARNING: Please remove 'uvicorn chromadb.app:app' from your command line arguments. This is now handled by the entrypoint script." + exec $(eval echo "$args") +else + echo "Starting 'uvicorn chromadb.app:app' with args: $(eval echo "$args")" + exec uvicorn chromadb.app:app $(eval echo "$args") +fi diff --git a/bin/reset.sh b/bin/reset.sh deleted file mode 100755 index 92fb04d8224..00000000000 --- a/bin/reset.sh +++ /dev/null @@ -1,13 +0,0 @@ - #!/usr/bin/env bash - -eval $(minikube -p chroma-test docker-env) - -docker build -t chroma-coordinator:latest -f go/coordinator/Dockerfile . - -kubectl delete deployment coordinator -n chroma - -# Apply the kubernetes manifests -kubectl apply -f k8s/deployment -kubectl apply -f k8s/crd -kubectl apply -f k8s/cr -kubectl apply -f k8s/test diff --git a/chromadb/__init__.py b/chromadb/__init__.py index 142ab78a05f..341aabe9ed8 100644 --- a/chromadb/__init__.py +++ b/chromadb/__init__.py @@ -43,7 +43,7 @@ __settings = Settings() -__version__ = "0.4.22" +__version__ = "0.4.24" # Workaround to deal with Colab's old sqlite3 version try: @@ -112,6 +112,10 @@ def EphemeralClient( settings = Settings() settings.is_persistent = False + # Make sure paramaters are the correct types -- users can pass anything. + tenant = str(tenant) + database = str(database) + return ClientCreator(settings=settings, tenant=tenant, database=database) @@ -135,12 +139,16 @@ def PersistentClient( settings.persist_directory = path settings.is_persistent = True + # Make sure paramaters are the correct types -- users can pass anything. + tenant = str(tenant) + database = str(database) + return ClientCreator(tenant=tenant, database=database, settings=settings) def HttpClient( host: str = "localhost", - port: str = "8000", + port: int = 8000, ssl: bool = False, headers: Optional[Dict[str, str]] = None, settings: Optional[Settings] = None, @@ -165,6 +173,13 @@ def HttpClient( if settings is None: settings = Settings() + # Make sure paramaters are the correct types -- users can pass anything. + host = str(host) + port = int(port) + ssl = bool(ssl) + tenant = str(tenant) + database = str(database) + settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI" if settings.chroma_server_host and settings.chroma_server_host != host: raise ValueError( @@ -189,7 +204,7 @@ def CloudClient( settings: Optional[Settings] = None, *, # Following arguments are keyword-only, intended for testing only. cloud_host: str = "api.trychroma.com", - cloud_port: str = "8000", + cloud_port: int = 8000, enable_ssl: bool = True, ) -> ClientAPI: """ @@ -217,6 +232,14 @@ def CloudClient( if settings is None: settings = Settings() + # Make sure paramaters are the correct types -- users can pass anything. + tenant = str(tenant) + database = str(database) + api_key = str(api_key) + cloud_host = str(cloud_host) + cloud_port = int(cloud_port) + enable_ssl = bool(enable_ssl) + settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI" settings.chroma_server_host = cloud_host settings.chroma_server_http_port = cloud_port @@ -242,9 +265,12 @@ def Client( tenant: The tenant to use for this client. Defaults to the default tenant. database: The database to use for this client. Defaults to the default database. - """ + # Make sure paramaters are the correct types -- users can pass anything. + tenant = str(tenant) + database = str(database) + return ClientCreator(tenant=tenant, database=database, settings=settings) diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py index d3d1a8a4e7e..d01028c734f 100644 --- a/chromadb/api/fastapi.py +++ b/chromadb/api/fastapi.py @@ -1,4 +1,4 @@ -import json +import orjson as json import logging from typing import Optional, cast, Tuple from typing import Sequence @@ -109,7 +109,7 @@ def __init__(self, system: System): self._api_url = FastAPI.resolve_url( chroma_server_host=str(system.settings.chroma_server_host), - chroma_server_http_port=int(str(system.settings.chroma_server_http_port)), + chroma_server_http_port=system.settings.chroma_server_http_port, chroma_server_ssl_enabled=system.settings.chroma_server_ssl_enabled, default_api_path=system.settings.chroma_server_api_default_path, ) @@ -138,6 +138,8 @@ def __init__(self, system: System): self._session = requests.Session() if self._header is not None: self._session.headers.update(self._header) + if self._settings.chroma_server_ssl_verify is not None: + self._session.verify = self._settings.chroma_server_ssl_verify @trace_method("FastAPI.heartbeat", OpenTelemetryGranularity.OPERATION) @override @@ -145,7 +147,7 @@ def heartbeat(self) -> int: """Returns the current server time in nanoseconds to check if the server is alive""" resp = self._session.get(self._api_url) raise_chroma_error(resp) - return int(resp.json()["nanosecond heartbeat"]) + return int(json.loads(resp.text)["nanosecond heartbeat"]) @trace_method("FastAPI.create_database", OpenTelemetryGranularity.OPERATION) @override @@ -175,7 +177,7 @@ def get_database( params={"tenant": tenant}, ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Database( id=resp_json["id"], name=resp_json["name"], tenant=resp_json["tenant"] ) @@ -196,7 +198,7 @@ def get_tenant(self, name: str) -> Tenant: self._api_url + "/tenants/" + name, ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Tenant(name=resp_json["name"]) @trace_method("FastAPI.list_collections", OpenTelemetryGranularity.OPERATION) @@ -219,7 +221,7 @@ def list_collections( }, ) raise_chroma_error(resp) - json_collections = resp.json() + json_collections = json.loads(resp.text) collections = [] for json_collection in json_collections: collections.append(Collection(self, **json_collection)) @@ -237,7 +239,7 @@ def count_collections( params={"tenant": tenant, "database": database}, ) raise_chroma_error(resp) - return cast(int, resp.json()) + return cast(int, json.loads(resp.text)) @trace_method("FastAPI.create_collection", OpenTelemetryGranularity.OPERATION) @override @@ -266,7 +268,7 @@ def create_collection( params={"tenant": tenant, "database": database}, ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Collection( client=self, id=resp_json["id"], @@ -300,7 +302,7 @@ def get_collection( self._api_url + "/collections/" + name if name else str(id), params=_params ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Collection( client=self, name=resp_json["name"], @@ -379,7 +381,7 @@ def _count( self._api_url + "/collections/" + str(collection_id) + "/count" ) raise_chroma_error(resp) - return cast(int, resp.json()) + return cast(int, json.loads(resp.text)) @trace_method("FastAPI._peek", OpenTelemetryGranularity.OPERATION) @override @@ -432,7 +434,7 @@ def _get( ) raise_chroma_error(resp) - body = resp.json() + body = json.loads(resp.text) return GetResult( ids=body["ids"], embeddings=body.get("embeddings", None), @@ -460,7 +462,7 @@ def _delete( ) raise_chroma_error(resp) - return cast(IDs, resp.json()) + return cast(IDs, json.loads(resp.text)) @trace_method("FastAPI._submit_batch", OpenTelemetryGranularity.ALL) def _submit_batch( @@ -584,7 +586,7 @@ def _query( ) raise_chroma_error(resp) - body = resp.json() + body = json.loads(resp.text) return QueryResult( ids=body["ids"], @@ -602,7 +604,7 @@ def reset(self) -> bool: """Resets the database""" resp = self._session.post(self._api_url + "/reset") raise_chroma_error(resp) - return cast(bool, resp.json()) + return cast(bool, json.loads(resp.text)) @trace_method("FastAPI.get_version", OpenTelemetryGranularity.OPERATION) @override @@ -610,7 +612,7 @@ def get_version(self) -> str: """Returns the version of the server""" resp = self._session.get(self._api_url + "/version") raise_chroma_error(resp) - return cast(str, resp.json()) + return cast(str, json.loads(resp.text)) @override def get_settings(self) -> Settings: @@ -624,7 +626,7 @@ def max_batch_size(self) -> int: if self._max_batch_size == -1: resp = self._session.get(self._api_url + "/pre-flight-checks") raise_chroma_error(resp) - self._max_batch_size = cast(int, resp.json()["max_batch_size"]) + self._max_batch_size = cast(int, json.loads(resp.text)["max_batch_size"]) return self._max_batch_size @@ -635,7 +637,7 @@ def raise_chroma_error(resp: requests.Response) -> None: chroma_error = None try: - body = resp.json() + body = json.loads(resp.text) if "error" in body: if body["error"] in errors.error_types: chroma_error = errors.error_types[body["error"]](body["message"]) diff --git a/chromadb/api/segment.py b/chromadb/api/segment.py index 72df138d9be..ce19f0a379c 100644 --- a/chromadb/api/segment.py +++ b/chromadb/api/segment.py @@ -1,6 +1,8 @@ from chromadb.api import ServerAPI from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System from chromadb.db.system import SysDB +from chromadb.quota import QuotaEnforcer, Resource +from chromadb.rate_limiting import rate_limit from chromadb.segment import SegmentManager, MetadataReader, VectorReader from chromadb.telemetry.opentelemetry import ( add_attributes_to_current_span, @@ -50,7 +52,6 @@ ) import chromadb.types as t - from typing import Any, Optional, Sequence, Generator, List, cast, Set, Dict from overrides import override from uuid import UUID, uuid4 @@ -58,7 +59,6 @@ import logging import re - logger = logging.getLogger(__name__) @@ -101,6 +101,7 @@ def __init__(self, system: System): self._settings = system.settings self._sysdb = self.require(SysDB) self._manager = self.require(SegmentManager) + self._quota = self.require(QuotaEnforcer) self._product_telemetry_client = self.require(ProductTelemetryClient) self._opentelemetry_client = self.require(OpenTelemetryClient) self._producer = self.require(Producer) @@ -110,6 +111,7 @@ def __init__(self, system: System): def heartbeat(self) -> int: return int(time.time_ns()) + @trace_method("SegmentAPI.create_database", OpenTelemetryGranularity.OPERATION) @override def create_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None: if len(name) < 3: @@ -121,10 +123,12 @@ def create_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None: tenant=tenant, ) + @trace_method("SegmentAPI.get_database", OpenTelemetryGranularity.OPERATION) @override def get_database(self, name: str, tenant: str = DEFAULT_TENANT) -> t.Database: return self._sysdb.get_database(name=name, tenant=tenant) + @trace_method("SegmentAPI.create_tenant", OpenTelemetryGranularity.OPERATION) @override def create_tenant(self, name: str) -> None: if len(name) < 3: @@ -134,6 +138,7 @@ def create_tenant(self, name: str) -> None: name=name, ) + @trace_method("SegmentAPI.get_tenant", OpenTelemetryGranularity.OPERATION) @override def get_tenant(self, name: str) -> t.Tenant: return self._sysdb.get_tenant(name=name) @@ -173,10 +178,13 @@ def create_collection( database=database, ) + # TODO: wrap sysdb call in try except and log error if it fails if created: segments = self._manager.create_segments(coll) for segment in segments: self._sysdb.create_segment(segment) + else: + logger.info(f"Collection {name} is not created.") # TODO: This event doesn't capture the get_or_create case appropriately self._product_telemetry_client.capture( @@ -346,6 +354,7 @@ def delete_collection( raise ValueError(f"Collection {name} does not exist.") @trace_method("SegmentAPI._add", OpenTelemetryGranularity.OPERATION) + @rate_limit(subject="collection_id", resource=Resource.ADD_PER_MINUTE) @override def _add( self, @@ -356,6 +365,7 @@ def _add( documents: Optional[Documents] = None, uris: Optional[URIs] = None, ) -> bool: + self._quota.static_check(metadatas, documents, embeddings, str(collection_id)) coll = self._get_collection(collection_id) self._manager.hint_use_collection(collection_id, t.Operation.ADD) validate_batch( @@ -366,7 +376,6 @@ def _add( for r in _records( t.Operation.ADD, ids=ids, - collection_id=collection_id, embeddings=embeddings, metadatas=metadatas, documents=documents, @@ -374,7 +383,7 @@ def _add( ): self._validate_embedding_record(coll, r) records_to_submit.append(r) - self._producer.submit_embeddings(coll["topic"], records_to_submit) + self._producer.submit_embeddings(collection_id, records_to_submit) self._product_telemetry_client.capture( CollectionAddEvent( @@ -398,6 +407,7 @@ def _update( documents: Optional[Documents] = None, uris: Optional[URIs] = None, ) -> bool: + self._quota.static_check(metadatas, documents, embeddings, str(collection_id)) coll = self._get_collection(collection_id) self._manager.hint_use_collection(collection_id, t.Operation.UPDATE) validate_batch( @@ -408,7 +418,6 @@ def _update( for r in _records( t.Operation.UPDATE, ids=ids, - collection_id=collection_id, embeddings=embeddings, metadatas=metadatas, documents=documents, @@ -416,7 +425,7 @@ def _update( ): self._validate_embedding_record(coll, r) records_to_submit.append(r) - self._producer.submit_embeddings(coll["topic"], records_to_submit) + self._producer.submit_embeddings(collection_id, records_to_submit) self._product_telemetry_client.capture( CollectionUpdateEvent( @@ -442,6 +451,7 @@ def _upsert( documents: Optional[Documents] = None, uris: Optional[URIs] = None, ) -> bool: + self._quota.static_check(metadatas, documents, embeddings, str(collection_id)) coll = self._get_collection(collection_id) self._manager.hint_use_collection(collection_id, t.Operation.UPSERT) validate_batch( @@ -452,7 +462,6 @@ def _upsert( for r in _records( t.Operation.UPSERT, ids=ids, - collection_id=collection_id, embeddings=embeddings, metadatas=metadatas, documents=documents, @@ -460,11 +469,12 @@ def _upsert( ): self._validate_embedding_record(coll, r) records_to_submit.append(r) - self._producer.submit_embeddings(coll["topic"], records_to_submit) + self._producer.submit_embeddings(collection_id, records_to_submit) return True @trace_method("SegmentAPI._get", OpenTelemetryGranularity.OPERATION) + @rate_limit(subject="collection_id", resource=Resource.GET_PER_MINUTE) @override def _get( self, @@ -621,12 +631,10 @@ def _delete( return [] records_to_submit = [] - for r in _records( - operation=t.Operation.DELETE, ids=ids_to_delete, collection_id=collection_id - ): + for r in _records(operation=t.Operation.DELETE, ids=ids_to_delete): self._validate_embedding_record(coll, r) records_to_submit.append(r) - self._producer.submit_embeddings(coll["topic"], records_to_submit) + self._producer.submit_embeddings(collection_id, records_to_submit) self._product_telemetry_client.capture( CollectionDeleteEvent( @@ -643,6 +651,7 @@ def _count(self, collection_id: UUID) -> int: return metadata_segment.count() @trace_method("SegmentAPI._query", OpenTelemetryGranularity.OPERATION) + @rate_limit(subject="collection_id", resource=Resource.QUERY_PER_MINUTE) @override def _query( self, @@ -791,7 +800,7 @@ def max_batch_size(self) -> int: # used for channel assignment in the distributed version of the system. @trace_method("SegmentAPI._validate_embedding_record", OpenTelemetryGranularity.ALL) def _validate_embedding_record( - self, collection: t.Collection, record: t.SubmitEmbeddingRecord + self, collection: t.Collection, record: t.OperationRecord ) -> None: """Validate the dimension of an embedding record before submitting it to the system.""" add_attributes_to_current_span({"collection_id": str(collection["id"])}) @@ -833,12 +842,11 @@ def _get_collection(self, collection_id: UUID) -> t.Collection: def _records( operation: t.Operation, ids: IDs, - collection_id: UUID, embeddings: Optional[Embeddings] = None, metadatas: Optional[Metadatas] = None, documents: Optional[Documents] = None, uris: Optional[URIs] = None, -) -> Generator[t.SubmitEmbeddingRecord, None, None]: +) -> Generator[t.OperationRecord, None, None]: """Convert parallel lists of embeddings, metadatas and documents to a sequence of SubmitEmbeddingRecords""" @@ -865,13 +873,12 @@ def _records( else: metadata = {"chroma:uri": uri} - record = t.SubmitEmbeddingRecord( + record = t.OperationRecord( id=id, embedding=embeddings[i] if embeddings else None, encoding=t.ScalarEncoding.FLOAT32, # Hardcode for now metadata=metadata, operation=operation, - collection_id=collection_id, ) yield record diff --git a/chromadb/api/types.py b/chromadb/api/types.py index 7781c422572..1b85dbc5fe5 100644 --- a/chromadb/api/types.py +++ b/chromadb/api/types.py @@ -20,7 +20,7 @@ # Re-export types from chromadb.types __all__ = ["Metadata", "Where", "WhereDocument", "UpdateCollectionMetadata"] - +META_KEY_CHROMA_DOCUMENT = "chroma:document" T = TypeVar("T") OneOrMany = Union[T, List[T]] @@ -223,9 +223,9 @@ def __call__(self, uris: URIs) -> L: def validate_ids(ids: IDs) -> IDs: """Validates ids to ensure it is a list of strings""" if not isinstance(ids, list): - raise ValueError(f"Expected IDs to be a list, got {ids}") + raise ValueError(f"Expected IDs to be a list, got {type(ids).__name__} as IDs") if len(ids) == 0: - raise ValueError(f"Expected IDs to be a non-empty list, got {ids}") + raise ValueError(f"Expected IDs to be a non-empty list, got {len(ids)} IDs") seen = set() dups = set() for id_ in ids: @@ -259,20 +259,28 @@ def validate_ids(ids: IDs) -> IDs: def validate_metadata(metadata: Metadata) -> Metadata: """Validates metadata to ensure it is a dictionary of strings to strings, ints, floats or bools""" if not isinstance(metadata, dict) and metadata is not None: - raise ValueError(f"Expected metadata to be a dict or None, got {metadata}") + raise ValueError( + f"Expected metadata to be a dict or None, got {type(metadata).__name__} as metadata" + ) if metadata is None: return metadata if len(metadata) == 0: - raise ValueError(f"Expected metadata to be a non-empty dict, got {metadata}") + raise ValueError( + f"Expected metadata to be a non-empty dict, got {len(metadata)} metadata attributes" + ) for key, value in metadata.items(): + if key == META_KEY_CHROMA_DOCUMENT: + raise ValueError( + f"Expected metadata to not contain the reserved key {META_KEY_CHROMA_DOCUMENT}" + ) if not isinstance(key, str): raise TypeError( - f"Expected metadata key to be a str, got {key} which is a {type(key)}" + f"Expected metadata key to be a str, got {key} which is a {type(key).__name__}" ) # isinstance(True, int) evaluates to True, so we need to check for bools separately if not isinstance(value, bool) and not isinstance(value, (str, int, float)): raise ValueError( - f"Expected metadata value to be a str, int, float or bool, got {value} which is a {type(value)}" + f"Expected metadata value to be a str, int, float or bool, got {value} which is a {type(value).__name__}" ) return metadata @@ -280,7 +288,9 @@ def validate_metadata(metadata: Metadata) -> Metadata: def validate_update_metadata(metadata: UpdateMetadata) -> UpdateMetadata: """Validates metadata to ensure it is a dictionary of strings to strings, ints, floats or bools""" if not isinstance(metadata, dict) and metadata is not None: - raise ValueError(f"Expected metadata to be a dict or None, got {metadata}") + raise ValueError( + f"Expected metadata to be a dict or None, got {type(metadata)}" + ) if metadata is None: return metadata if len(metadata) == 0: @@ -467,16 +477,23 @@ def validate_n_results(n_results: int) -> int: def validate_embeddings(embeddings: Embeddings) -> Embeddings: """Validates embeddings to ensure it is a list of list of ints, or floats""" if not isinstance(embeddings, list): - raise ValueError(f"Expected embeddings to be a list, got {embeddings}") + raise ValueError( + f"Expected embeddings to be a list, got {type(embeddings).__name__}" + ) if len(embeddings) == 0: raise ValueError( - f"Expected embeddings to be a list with at least one item, got {embeddings}" + f"Expected embeddings to be a list with at least one item, got {len(embeddings)} embeddings" ) if not all([isinstance(e, list) for e in embeddings]): raise ValueError( - f"Expected each embedding in the embeddings to be a list, got {embeddings}" + "Expected each embedding in the embeddings to be a list, got " + f"{list(set([type(e).__name__ for e in embeddings]))}" ) - for embedding in embeddings: + for i, embedding in enumerate(embeddings): + if len(embedding) == 0: + raise ValueError( + f"Expected each embedding in the embeddings to be a non-empty list, got empty embedding at pos {i}" + ) if not all( [ isinstance(value, (int, float)) and not isinstance(value, bool) @@ -484,7 +501,8 @@ def validate_embeddings(embeddings: Embeddings) -> Embeddings: ] ): raise ValueError( - f"Expected each value in the embedding to be a int or float, got {embeddings}" + "Expected each value in the embedding to be a int or float, got an embedding with " + f"{list(set([type(value).__name__ for value in embedding]))} - {embedding}" ) return embeddings diff --git a/chromadb/auth/fastapi.py b/chromadb/auth/fastapi.py index 1f9d3c900c3..85c5b803135 100644 --- a/chromadb/auth/fastapi.py +++ b/chromadb/auth/fastapi.py @@ -1,4 +1,3 @@ -import chromadb from contextvars import ContextVar from functools import wraps import logging @@ -9,6 +8,7 @@ from starlette.responses import Response from starlette.types import ASGIApp +import chromadb from chromadb.config import DEFAULT_TENANT, System from chromadb.auth import ( AuthorizationContext, @@ -28,7 +28,7 @@ ) from chromadb.auth.registry import resolve_provider from chromadb.errors import AuthorizationError -from chromadb.server.fastapi.utils import fastapi_json_response +from chromadb.utils.fastapi import fastapi_json_response from chromadb.telemetry.opentelemetry import ( OpenTelemetryGranularity, trace_method, @@ -117,7 +117,7 @@ def instrument_server(self, app: ASGIApp) -> None: raise NotImplementedError("Not implemented yet") -class FastAPIChromaAuthMiddlewareWrapper(BaseHTTPMiddleware): # type: ignore +class FastAPIChromaAuthMiddlewareWrapper(BaseHTTPMiddleware): def __init__( self, app: ASGIApp, auth_middleware: FastAPIChromaAuthMiddleware ) -> None: @@ -302,7 +302,7 @@ def instrument_server(self, app: ASGIApp) -> None: raise NotImplementedError("Not implemented yet") -class FastAPIChromaAuthzMiddlewareWrapper(BaseHTTPMiddleware): # type: ignore +class FastAPIChromaAuthzMiddlewareWrapper(BaseHTTPMiddleware): def __init__( self, app: ASGIApp, authz_middleware: FastAPIChromaAuthzMiddleware ) -> None: diff --git a/chromadb/auth/fastapi_utils.py b/chromadb/auth/fastapi_utils.py index 2612bf6716f..9dc652feeb9 100644 --- a/chromadb/auth/fastapi_utils.py +++ b/chromadb/auth/fastapi_utils.py @@ -1,6 +1,6 @@ from functools import partial from typing import Any, Callable, Dict, Optional, Sequence, cast -from chromadb.server.fastapi.utils import string_to_uuid +from chromadb.utils.fastapi import string_to_uuid from chromadb.api import ServerAPI from chromadb.auth import AuthzResourceTypes diff --git a/chromadb/config.py b/chromadb/config.py index 59bea5ee0e4..7fd8e6d8981 100644 --- a/chromadb/config.py +++ b/chromadb/config.py @@ -4,7 +4,7 @@ import os from abc import ABC from graphlib import TopologicalSorter -from typing import Optional, List, Any, Dict, Set, Iterable +from typing import Optional, List, Any, Dict, Set, Iterable, Union from typing import Type, TypeVar, cast from overrides import EnforceOverrides @@ -70,11 +70,12 @@ "chromadb.telemetry.product.ProductTelemetryClient": "chroma_product_telemetry_impl", "chromadb.ingest.Producer": "chroma_producer_impl", "chromadb.ingest.Consumer": "chroma_consumer_impl", - "chromadb.ingest.CollectionAssignmentPolicy": "chroma_collection_assignment_policy_impl", # noqa + "chromadb.quota.QuotaProvider": "chroma_quota_provider_impl", "chromadb.db.system.SysDB": "chroma_sysdb_impl", "chromadb.segment.SegmentManager": "chroma_segment_manager_impl", "chromadb.segment.distributed.SegmentDirectory": "chroma_segment_directory_impl", "chromadb.segment.distributed.MemberlistProvider": "chroma_memberlist_provider_impl", + "chromadb.rate_limiting.RateLimitingProvider": "chroma_rate_limiting_provider_impl", } DEFAULT_TENANT = "default_tenant" @@ -84,9 +85,12 @@ class Settings(BaseSettings): # type: ignore environment: str = "" - # Legacy config has to be kept around because pydantic will error + # Legacy config that has to be kept around because pydantic will error # on nonexisting keys chroma_db_impl: Optional[str] = None + chroma_collection_assignment_policy_impl: str = ( + "chromadb.ingest.impl.simple_policy.SimpleAssignmentPolicy" + ) # Can be "chromadb.api.segment.SegmentAPI" or "chromadb.api.fastapi.FastAPI" chroma_api_impl: str = "chromadb.api.segment.SegmentAPI" chroma_product_telemetry_impl: str = "chromadb.telemetry.product.posthog.Posthog" @@ -101,27 +105,35 @@ class Settings(BaseSettings): # type: ignore "chromadb.segment.impl.manager.local.LocalSegmentManager" ) + chroma_quota_provider_impl: Optional[str] = None + chroma_rate_limiting_provider_impl: Optional[str] = None + # Distributed architecture specific components chroma_segment_directory_impl: str = "chromadb.segment.impl.distributed.segment_directory.RendezvousHashSegmentDirectory" chroma_memberlist_provider_impl: str = "chromadb.segment.impl.distributed.segment_directory.CustomResourceMemberlistProvider" - chroma_collection_assignment_policy_impl: str = ( - "chromadb.ingest.impl.simple_policy.SimpleAssignmentPolicy" - ) - worker_memberlist_name: str = "worker-memberlist" + worker_memberlist_name: str = "query-service-memberlist" chroma_coordinator_host = "localhost" + chroma_logservice_host = "localhost" + chroma_logservice_port = 50052 + tenant_id: str = "default" topic_namespace: str = "default" is_persistent: bool = False persist_directory: str = "./chroma" + chroma_memory_limit_bytes: int = 0 + chroma_segment_cache_policy: Optional[str] = None + chroma_server_host: Optional[str] = None chroma_server_headers: Optional[Dict[str, str]] = None - chroma_server_http_port: Optional[str] = None + chroma_server_http_port: Optional[int] = None chroma_server_ssl_enabled: Optional[bool] = False + # the below config value is only applicable to Chroma HTTP clients + chroma_server_ssl_verify: Optional[Union[bool, str]] = None chroma_server_api_default_path: Optional[str] = "/api/v1" - chroma_server_grpc_port: Optional[str] = None + chroma_server_grpc_port: Optional[int] = None # eg ["http://localhost:3000"] chroma_server_cors_allow_origins: List[str] = [] @@ -133,10 +145,6 @@ def empty_str_to_none(cls, v: str) -> Optional[str]: chroma_server_nofile: Optional[int] = None - pulsar_broker_url: Optional[str] = None - pulsar_admin_port: Optional[str] = "8080" - pulsar_broker_port: Optional[str] = "6650" - chroma_server_auth_provider: Optional[str] = None @validator("chroma_server_auth_provider", pre=True, always=True, allow_reuse=True) @@ -313,6 +321,18 @@ def __init__(self, settings: Settings): if settings[key] is not None: raise ValueError(LEGACY_ERROR) + if ( + settings["chroma_segment_cache_policy"] is not None + and settings["chroma_segment_cache_policy"] != "LRU" + ): + logger.error( + "Failed to set chroma_segment_cache_policy: Only LRU is available." + ) + if settings["chroma_memory_limit_bytes"] == 0: + logger.error( + "Failed to set chroma_segment_cache_policy: chroma_memory_limit_bytes is require." + ) + # Apply the nofile limit if set if settings["chroma_server_nofile"] is not None: if platform.system() != "Windows": diff --git a/chromadb/db/impl/grpc/client.py b/chromadb/db/impl/grpc/client.py index 32b3b2da164..ebdf59445e8 100644 --- a/chromadb/db/impl/grpc/client.py +++ b/chromadb/db/impl/grpc/client.py @@ -1,7 +1,8 @@ +import logging from typing import List, Optional, Sequence, Tuple, Union, cast from uuid import UUID from overrides import overrides -from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, System +from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, System, logger from chromadb.db.base import NotFoundError, UniqueConstraintError from chromadb.db.system import SysDB from chromadb.proto.convert import ( @@ -27,6 +28,7 @@ UpdateSegmentRequest, ) from chromadb.proto.coordinator_pb2_grpc import SysDBStub +from chromadb.telemetry.opentelemetry.grpc import OtelInterceptor from chromadb.types import ( Collection, Database, @@ -64,6 +66,8 @@ def start(self) -> None: self._channel = grpc.insecure_channel( f"{self._coordinator_url}:{self._coordinator_port}" ) + interceptors = [OtelInterceptor()] + self._channel = grpc.intercept_channel(self._channel, *interceptors) self._sys_db_stub = SysDBStub(self._channel) # type: ignore return super().start() @@ -140,14 +144,12 @@ def get_segments( id: Optional[UUID] = None, type: Optional[str] = None, scope: Optional[SegmentScope] = None, - topic: Optional[str] = None, collection: Optional[UUID] = None, ) -> Sequence[Segment]: request = GetSegmentsRequest( id=id.hex if id else None, type=type, scope=to_proto_segment_scope(scope) if scope else None, - topic=topic, collection=collection.hex if collection else None, ) response = self._sys_db_stub.GetSegments(request) @@ -161,14 +163,9 @@ def get_segments( def update_segment( self, id: UUID, - topic: OptionalArgument[Optional[str]] = Unspecified(), collection: OptionalArgument[Optional[UUID]] = Unspecified(), metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(), ) -> None: - write_topic = None - if topic != Unspecified(): - write_topic = cast(Union[str, None], topic) - write_collection = None if collection != Unspecified(): write_collection = cast(Union[UUID, None], collection) @@ -179,17 +176,12 @@ def update_segment( request = UpdateSegmentRequest( id=id.hex, - topic=write_topic, collection=write_collection.hex if write_collection else None, metadata=to_proto_update_metadata(write_metadata) if write_metadata else None, ) - if topic is None: - request.ClearField("topic") - request.reset_topic = True - if collection is None: request.ClearField("collection") request.reset_collection = True @@ -221,10 +213,13 @@ def create_collection( database=database, ) response = self._sys_db_stub.CreateCollection(request) + # TODO: this needs to be changed to try, catch instead of checking the status code + if response.status.code != 200: + logger.info(f"failed to create collection, response: {response}") if response.status.code == 409: raise UniqueConstraintError() collection = from_proto_collection(response.collection) - return collection, response.created + return collection, response.status.code == 200 @overrides def delete_collection( @@ -236,6 +231,7 @@ def delete_collection( database=database, ) response = self._sys_db_stub.DeleteCollection(request) + logging.debug(f"delete_collection response: {response}") if response.status.code == 404: raise NotFoundError() @@ -243,7 +239,6 @@ def delete_collection( def get_collections( self, id: Optional[UUID] = None, - topic: Optional[str] = None, name: Optional[str] = None, tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, @@ -253,7 +248,6 @@ def get_collections( # TODO: implement limit and offset in the gRPC service request = GetCollectionsRequest( id=id.hex if id else None, - topic=topic, name=name, tenant=tenant, database=database, @@ -268,15 +262,10 @@ def get_collections( def update_collection( self, id: UUID, - topic: OptionalArgument[str] = Unspecified(), name: OptionalArgument[str] = Unspecified(), dimension: OptionalArgument[Optional[int]] = Unspecified(), metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(), ) -> None: - write_topic = None - if topic != Unspecified(): - write_topic = cast(str, topic) - write_name = None if name != Unspecified(): write_name = cast(str, name) @@ -291,7 +280,6 @@ def update_collection( request = UpdateCollectionRequest( id=id.hex, - topic=write_topic, name=write_name, dimension=write_dimension, metadata=to_proto_update_metadata(write_metadata) diff --git a/chromadb/db/impl/grpc/server.py b/chromadb/db/impl/grpc/server.py index 257aa80f0e7..cdec22eff0d 100644 --- a/chromadb/db/impl/grpc/server.py +++ b/chromadb/db/impl/grpc/server.py @@ -2,7 +2,6 @@ from typing import Any, Dict, cast from uuid import UUID from overrides import overrides -from chromadb.ingest import CollectionAssignmentPolicy from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Component, System from chromadb.proto.convert import ( from_proto_metadata, @@ -17,9 +16,15 @@ CreateCollectionRequest, CreateCollectionResponse, CreateDatabaseRequest, + CreateDatabaseResponse, CreateSegmentRequest, + CreateSegmentResponse, + CreateTenantRequest, + CreateTenantResponse, DeleteCollectionRequest, + DeleteCollectionResponse, DeleteSegmentRequest, + DeleteSegmentResponse, GetCollectionsRequest, GetCollectionsResponse, GetDatabaseRequest, @@ -28,8 +33,11 @@ GetSegmentsResponse, GetTenantRequest, GetTenantResponse, + ResetStateResponse, UpdateCollectionRequest, + UpdateCollectionResponse, UpdateSegmentRequest, + UpdateSegmentResponse, ) from chromadb.proto.coordinator_pb2_grpc import ( SysDBServicer, @@ -46,7 +54,6 @@ class GrpcMockSysDB(SysDBServicer, Component): _server: grpc.Server _server_port: int - _assignment_policy: CollectionAssignmentPolicy _segments: Dict[str, Segment] = {} _tenants_to_databases_to_collections: Dict[ str, Dict[str, Dict[str, Collection]] @@ -55,7 +62,6 @@ class GrpcMockSysDB(SysDBServicer, Component): def __init__(self, system: System): self._server_port = system.settings.require("chroma_server_grpc_port") - self._assignment_policy = system.instance(CollectionAssignmentPolicy) return super().__init__(system) @overrides @@ -68,7 +74,7 @@ def start(self) -> None: @overrides def stop(self) -> None: - self._server.stop(0) + self._server.stop(None) return super().stop() @overrides @@ -85,22 +91,22 @@ def reset_state(self) -> None: @overrides(check_signature=False) def CreateDatabase( self, request: CreateDatabaseRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> CreateDatabaseResponse: tenant = request.tenant database = request.name if tenant not in self._tenants_to_databases_to_collections: - return proto.ChromaResponse( + return CreateDatabaseResponse( status=proto.Status(code=404, reason=f"Tenant {tenant} not found") ) if database in self._tenants_to_databases_to_collections[tenant]: - return proto.ChromaResponse( + return CreateDatabaseResponse( status=proto.Status( code=409, reason=f"Database {database} already exists" ) ) self._tenants_to_databases_to_collections[tenant][database] = {} self._tenants_to_database_to_id[tenant][database] = UUID(hex=request.id) - return proto.ChromaResponse(status=proto.Status(code=200)) + return CreateDatabaseResponse(status=proto.Status(code=200)) @overrides(check_signature=False) def GetDatabase( @@ -124,16 +130,16 @@ def GetDatabase( @overrides(check_signature=False) def CreateTenant( - self, request: CreateDatabaseRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + self, request: CreateTenantRequest, context: grpc.ServicerContext + ) -> CreateTenantResponse: tenant = request.name if tenant in self._tenants_to_databases_to_collections: - return proto.ChromaResponse( + return CreateTenantResponse( status=proto.Status(code=409, reason=f"Tenant {tenant} already exists") ) self._tenants_to_databases_to_collections[tenant] = {} self._tenants_to_database_to_id[tenant] = {} - return proto.ChromaResponse(status=proto.Status(code=200)) + return CreateTenantResponse(status=proto.Status(code=200)) @overrides(check_signature=False) def GetTenant( @@ -155,29 +161,29 @@ def GetTenant( @overrides(check_signature=False) def CreateSegment( self, request: CreateSegmentRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> CreateSegmentResponse: segment = from_proto_segment(request.segment) if segment["id"].hex in self._segments: - return proto.ChromaResponse( + return CreateSegmentResponse( status=proto.Status( code=409, reason=f"Segment {segment['id']} already exists" ) ) self._segments[segment["id"].hex] = segment - return proto.ChromaResponse( + return CreateSegmentResponse( status=proto.Status(code=200) ) # TODO: how are these codes used? Need to determine the standards for the code and reason. @overrides(check_signature=False) def DeleteSegment( self, request: DeleteSegmentRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> DeleteSegmentResponse: id_to_delete = request.id if id_to_delete in self._segments: del self._segments[id_to_delete] - return proto.ChromaResponse(status=proto.Status(code=200)) + return DeleteSegmentResponse(status=proto.Status(code=200)) else: - return proto.ChromaResponse( + return DeleteSegmentResponse( status=proto.Status( code=404, reason=f"Segment {id_to_delete} not found" ) @@ -194,7 +200,6 @@ def GetSegments( if request.HasField("scope") else None ) - target_topic = request.topic if request.HasField("topic") else None target_collection = ( UUID(hex=request.collection) if request.HasField("collection") else None ) @@ -207,8 +212,6 @@ def GetSegments( continue if target_scope and segment["scope"] != target_scope: continue - if target_topic and segment["topic"] != target_topic: - continue if target_collection and segment["collection"] != target_collection: continue found_segments.append(segment) @@ -219,20 +222,16 @@ def GetSegments( @overrides(check_signature=False) def UpdateSegment( self, request: UpdateSegmentRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> UpdateSegmentResponse: id_to_update = UUID(request.id) if id_to_update.hex not in self._segments: - return proto.ChromaResponse( + return UpdateSegmentResponse( status=proto.Status( code=404, reason=f"Segment {id_to_update} not found" ) ) else: segment = self._segments[id_to_update.hex] - if request.HasField("topic"): - segment["topic"] = request.topic - if request.HasField("reset_topic") and request.reset_topic: - segment["topic"] = None if request.HasField("collection"): segment["collection"] = UUID(hex=request.collection) if request.HasField("reset_collection") and request.reset_collection: @@ -244,7 +243,7 @@ def UpdateSegment( self._merge_metadata(target, request.metadata) if request.HasField("reset_metadata") and request.reset_metadata: segment["metadata"] = {} - return proto.ChromaResponse(status=proto.Status(code=200)) + return UpdateSegmentResponse(status=proto.Status(code=200)) @overrides(check_signature=False) def CreateCollection( @@ -317,7 +316,6 @@ def CreateCollection( name=request.name, metadata=from_proto_metadata(request.metadata), dimension=request.dimension, - topic=self._assignment_policy.assign_collection(id), database=database, tenant=tenant, ) @@ -331,24 +329,24 @@ def CreateCollection( @overrides(check_signature=False) def DeleteCollection( self, request: DeleteCollectionRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> DeleteCollectionResponse: collection_id = request.id tenant = request.tenant database = request.database if tenant not in self._tenants_to_databases_to_collections: - return proto.ChromaResponse( + return DeleteCollectionResponse( status=proto.Status(code=404, reason=f"Tenant {tenant} not found") ) if database not in self._tenants_to_databases_to_collections[tenant]: - return proto.ChromaResponse( + return DeleteCollectionResponse( status=proto.Status(code=404, reason=f"Database {database} not found") ) collections = self._tenants_to_databases_to_collections[tenant][database] if collection_id in collections: del collections[collection_id] - return proto.ChromaResponse(status=proto.Status(code=200)) + return DeleteCollectionResponse(status=proto.Status(code=200)) else: - return proto.ChromaResponse( + return DeleteCollectionResponse( status=proto.Status( code=404, reason=f"Collection {collection_id} not found" ) @@ -359,7 +357,6 @@ def GetCollections( self, request: GetCollectionsRequest, context: grpc.ServicerContext ) -> GetCollectionsResponse: target_id = UUID(hex=request.id) if request.HasField("id") else None - target_topic = request.topic if request.HasField("topic") else None target_name = request.name if request.HasField("name") else None tenant = request.tenant @@ -378,8 +375,6 @@ def GetCollections( for collection in collections.values(): if target_id and collection["id"] != target_id: continue - if target_topic and collection["topic"] != target_topic: - continue if target_name and collection["name"] != target_name: continue found_collections.append(collection) @@ -392,7 +387,7 @@ def GetCollections( @overrides(check_signature=False) def UpdateCollection( self, request: UpdateCollectionRequest, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> UpdateCollectionResponse: id_to_update = UUID(request.id) # Find the collection with this id collections = {} @@ -402,15 +397,13 @@ def UpdateCollection( collections = maybe_collections if id_to_update.hex not in collections: - return proto.ChromaResponse( + return UpdateCollectionResponse( status=proto.Status( code=404, reason=f"Collection {id_to_update} not found" ) ) else: collection = collections[id_to_update.hex] - if request.HasField("topic"): - collection["topic"] = request.topic if request.HasField("name"): collection["name"] = request.name if request.HasField("dimension"): @@ -433,14 +426,14 @@ def UpdateCollection( if request.reset_metadata: collection["metadata"] = {} - return proto.ChromaResponse(status=proto.Status(code=200)) + return UpdateCollectionResponse(status=proto.Status(code=200)) @overrides(check_signature=False) def ResetState( self, request: Empty, context: grpc.ServicerContext - ) -> proto.ChromaResponse: + ) -> ResetStateResponse: self.reset_state() - return proto.ChromaResponse(status=proto.Status(code=200)) + return ResetStateResponse(status=proto.Status(code=200)) def _merge_metadata(self, target: Metadata, source: proto.UpdateMetadata) -> None: target_metadata = cast(Dict[str, Any], target) diff --git a/chromadb/db/migrations.py b/chromadb/db/migrations.py index 951cb762c36..97ef029092a 100644 --- a/chromadb/db/migrations.py +++ b/chromadb/db/migrations.py @@ -1,3 +1,4 @@ +import sys from typing import Sequence from typing_extensions import TypedDict, NotRequired from importlib_resources.abc import Traversable @@ -253,7 +254,7 @@ def _read_migration_file(file: MigrationFile, hash_alg: str) -> Migration: sql = file["path"].read_text() if hash_alg == "md5": - hash = hashlib.md5(sql.encode("utf-8")).hexdigest() + hash = hashlib.md5(sql.encode("utf-8"), usedforsecurity=False).hexdigest() if sys.version_info >= (3, 9) else hashlib.md5(sql.encode("utf-8")).hexdigest() elif hash_alg == "sha256": hash = hashlib.sha256(sql.encode("utf-8")).hexdigest() else: diff --git a/chromadb/db/mixins/embeddings_queue.py b/chromadb/db/mixins/embeddings_queue.py index b5d745b9286..913e6dc347d 100644 --- a/chromadb/db/mixins/embeddings_queue.py +++ b/chromadb/db/mixins/embeddings_queue.py @@ -1,16 +1,17 @@ +import json from chromadb.db.base import SqlDB, ParameterValue, get_sql from chromadb.ingest import ( Producer, Consumer, - encode_vector, - decode_vector, ConsumerCallbackFn, + decode_vector, + encode_vector, ) from chromadb.types import ( - SubmitEmbeddingRecord, - EmbeddingRecord, - SeqId, + OperationRecord, + LogRecord, ScalarEncoding, + SeqId, Operation, ) from chromadb.config import System @@ -21,12 +22,13 @@ ) from overrides import override from collections import defaultdict -from typing import Sequence, Tuple, Optional, Dict, Set, cast +from typing import Sequence, Optional, Dict, Set, Tuple, cast from uuid import UUID from pypika import Table, functions import uuid -import json import logging +from chromadb.ingest.impl.utils import create_topic_name + logger = logging.getLogger(__name__) @@ -78,6 +80,8 @@ def __init__( _subscriptions: Dict[str, Set[Subscription]] _max_batch_size: Optional[int] + _tenant: str + _topic_namespace: str # How many variables are in the insert statement for a single record VARIABLES_PER_RECORD = 6 @@ -85,6 +89,8 @@ def __init__(self, system: System): self._subscriptions = defaultdict(set) self._max_batch_size = None self._opentelemetry_client = system.require(OpenTelemetryClient) + self._tenant = system.settings.require("tenant_id") + self._topic_namespace = system.settings.require("topic_namespace") super().__init__(system) @trace_method("SqlEmbeddingsQueue.reset_state", OpenTelemetryGranularity.ALL) @@ -93,14 +99,12 @@ def reset_state(self) -> None: super().reset_state() self._subscriptions = defaultdict(set) - @override - def create_topic(self, topic_name: str) -> None: - # Topic creation is implicit for this impl - pass - @trace_method("SqlEmbeddingsQueue.delete_topic", OpenTelemetryGranularity.ALL) @override - def delete_topic(self, topic_name: str) -> None: + def delete_log(self, collection_id: UUID) -> None: + topic_name = create_topic_name( + self._tenant, self._topic_namespace, collection_id + ) t = Table("embeddings_queue") q = ( self.querybuilder() @@ -115,17 +119,17 @@ def delete_topic(self, topic_name: str) -> None: @trace_method("SqlEmbeddingsQueue.submit_embedding", OpenTelemetryGranularity.ALL) @override def submit_embedding( - self, topic_name: str, embedding: SubmitEmbeddingRecord + self, collection_id: UUID, embedding: OperationRecord ) -> SeqId: if not self._running: raise RuntimeError("Component not running") - return self.submit_embeddings(topic_name, [embedding])[0] + return self.submit_embeddings(collection_id, [embedding])[0] @trace_method("SqlEmbeddingsQueue.submit_embeddings", OpenTelemetryGranularity.ALL) @override def submit_embeddings( - self, topic_name: str, embeddings: Sequence[SubmitEmbeddingRecord] + self, collection_id: UUID, embeddings: Sequence[OperationRecord] ) -> Sequence[SeqId]: if not self._running: raise RuntimeError("Component not running") @@ -142,6 +146,10 @@ def submit_embeddings( """ ) + topic_name = create_topic_name( + self._tenant, self._topic_namespace, collection_id + ) + t = Table("embeddings_queue") insert = ( self.querybuilder() @@ -180,13 +188,15 @@ def submit_embeddings( submit_embedding_record = embeddings[id_to_idx[id]] # We allow notifying consumers out of order relative to one call to # submit_embeddings so we do not reorder the records before submitting them - embedding_record = EmbeddingRecord( - id=id, - seq_id=seq_id, - embedding=submit_embedding_record["embedding"], - encoding=submit_embedding_record["encoding"], - metadata=submit_embedding_record["metadata"], - operation=submit_embedding_record["operation"], + embedding_record = LogRecord( + log_offset=seq_id, + operation_record=OperationRecord( + id=id, + embedding=submit_embedding_record["embedding"], + encoding=submit_embedding_record["encoding"], + metadata=submit_embedding_record["metadata"], + operation=submit_embedding_record["operation"], + ), ) embedding_records.append(embedding_record) self._notify_all(topic_name, embedding_records) @@ -196,7 +206,7 @@ def submit_embeddings( @override def subscribe( self, - topic_name: str, + collection_id: UUID, consume_fn: ConsumerCallbackFn, start: Optional[SeqId] = None, end: Optional[SeqId] = None, @@ -205,6 +215,10 @@ def subscribe( if not self._running: raise RuntimeError("Component not running") + topic_name = create_topic_name( + self._tenant, self._topic_namespace, collection_id + ) + subscription_id = id or uuid.uuid4() start, end = self._validate_range(start, end) @@ -265,7 +279,7 @@ def max_batch_size(self) -> int: OpenTelemetryGranularity.ALL, ) def _prepare_vector_encoding_metadata( - self, embedding: SubmitEmbeddingRecord + self, embedding: OperationRecord ) -> Tuple[Optional[bytes], Optional[str], Optional[str]]: if embedding["embedding"]: encoding_type = cast(ScalarEncoding, embedding["encoding"]) @@ -305,13 +319,15 @@ def _backfill(self, subscription: Subscription) -> None: self._notify_one( subscription, [ - EmbeddingRecord( - seq_id=row[0], - operation=_operation_codes_inv[row[1]], - id=row[2], - embedding=vector, - encoding=encoding, - metadata=json.loads(row[5]) if row[5] else None, + LogRecord( + log_offset=row[0], + operation_record=OperationRecord( + operation=_operation_codes_inv[row[1]], + id=row[2], + embedding=vector, + encoding=encoding, + metadata=json.loads(row[5]) if row[5] else None, + ), ) ], ) @@ -340,24 +356,22 @@ def _next_seq_id(self) -> int: return int(cur.fetchone()[0]) + 1 @trace_method("SqlEmbeddingsQueue._notify_all", OpenTelemetryGranularity.ALL) - def _notify_all(self, topic: str, embeddings: Sequence[EmbeddingRecord]) -> None: + def _notify_all(self, topic: str, embeddings: Sequence[LogRecord]) -> None: """Send a notification to each subscriber of the given topic.""" if self._running: for sub in self._subscriptions[topic]: self._notify_one(sub, embeddings) @trace_method("SqlEmbeddingsQueue._notify_one", OpenTelemetryGranularity.ALL) - def _notify_one( - self, sub: Subscription, embeddings: Sequence[EmbeddingRecord] - ) -> None: + def _notify_one(self, sub: Subscription, embeddings: Sequence[LogRecord]) -> None: """Send a notification to a single subscriber.""" # Filter out any embeddings that are not in the subscription range should_unsubscribe = False filtered_embeddings = [] for embedding in embeddings: - if embedding["seq_id"] <= sub.start: + if embedding["log_offset"] <= sub.start: continue - if embedding["seq_id"] > sub.end: + if embedding["log_offset"] > sub.end: should_unsubscribe = True break filtered_embeddings.append(embedding) diff --git a/chromadb/db/mixins/sysdb.py b/chromadb/db/mixins/sysdb.py index 7373aabd0f3..dbab7310e91 100644 --- a/chromadb/db/mixins/sysdb.py +++ b/chromadb/db/mixins/sysdb.py @@ -20,7 +20,7 @@ OpenTelemetryGranularity, trace_method, ) -from chromadb.ingest import CollectionAssignmentPolicy, Producer +from chromadb.ingest import Producer from chromadb.types import ( Database, OptionalArgument, @@ -35,13 +35,11 @@ class SqlSysDB(SqlDB, SysDB): - _assignment_policy: CollectionAssignmentPolicy - # Used only to delete topics on collection deletion. + # Used only to delete log streams on collection deletion. # TODO: refactor to remove this dependency into a separate interface _producer: Producer def __init__(self, system: System): - self._assignment_policy = system.instance(CollectionAssignmentPolicy) super().__init__(system) self._opentelemetry_client = system.require(OpenTelemetryClient) @@ -143,7 +141,6 @@ def create_segment(self, segment: Segment) -> None: "segment_id": str(segment["id"]), "segment_type": segment["type"], "segment_scope": segment["scope"].value, - "segment_topic": str(segment["topic"]), "collection": str(segment["collection"]), } ) @@ -156,14 +153,12 @@ def create_segment(self, segment: Segment) -> None: segments.id, segments.type, segments.scope, - segments.topic, segments.collection, ) .insert( ParameterValue(self.uuid_to_db(segment["id"])), ParameterValue(segment["type"]), ParameterValue(segment["scope"].value), - ParameterValue(segment["topic"]), ParameterValue(self.uuid_to_db(segment["collection"])), ) ) @@ -224,10 +219,8 @@ def create_collection( else: raise UniqueConstraintError(f"Collection {name} already exists") - topic = self._assignment_policy.assign_collection(id) collection = Collection( id=id, - topic=topic, name=name, metadata=metadata, dimension=dimension, @@ -244,14 +237,12 @@ def create_collection( .into(collections) .columns( collections.id, - collections.topic, collections.name, collections.dimension, collections.database_id, ) .insert( ParameterValue(self.uuid_to_db(collection["id"])), - ParameterValue(collection["topic"]), ParameterValue(collection["name"]), ParameterValue(collection["dimension"]), # Get the database id for the database with the given name and tenant @@ -287,7 +278,6 @@ def get_segments( id: Optional[UUID] = None, type: Optional[str] = None, scope: Optional[SegmentScope] = None, - topic: Optional[str] = None, collection: Optional[UUID] = None, ) -> Sequence[Segment]: add_attributes_to_current_span( @@ -295,7 +285,6 @@ def get_segments( "segment_id": str(id), "segment_type": type if type else "", "segment_scope": scope.value if scope else "", - "segment_topic": topic if topic else "", "collection": str(collection), } ) @@ -308,7 +297,6 @@ def get_segments( segments_t.id, segments_t.type, segments_t.scope, - segments_t.topic, segments_t.collection, metadata_t.key, metadata_t.str_value, @@ -325,8 +313,6 @@ def get_segments( q = q.where(segments_t.type == ParameterValue(type)) if scope: q = q.where(segments_t.scope == ParameterValue(scope.value)) - if topic: - q = q.where(segments_t.topic == ParameterValue(topic)) if collection: q = q.where( segments_t.collection == ParameterValue(self.uuid_to_db(collection)) @@ -342,15 +328,13 @@ def get_segments( rows = list(segment_rows) type = str(rows[0][1]) scope = SegmentScope(str(rows[0][2])) - topic = str(rows[0][3]) if rows[0][3] else None - collection = self.uuid_from_db(rows[0][4]) if rows[0][4] else None + collection = self.uuid_from_db(rows[0][3]) if rows[0][3] else None metadata = self._metadata_from_rows(rows) segments.append( Segment( id=cast(UUID, id), type=type, scope=scope, - topic=topic, collection=collection, metadata=metadata, ) @@ -363,7 +347,6 @@ def get_segments( def get_collections( self, id: Optional[UUID] = None, - topic: Optional[str] = None, name: Optional[str] = None, tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, @@ -380,7 +363,6 @@ def get_collections( add_attributes_to_current_span( { "collection_id": str(id), - "collection_topic": topic if topic else "", "collection_name": name if name else "", } ) @@ -394,7 +376,6 @@ def get_collections( .select( collections_t.id, collections_t.name, - collections_t.topic, collections_t.dimension, databases_t.name, databases_t.tenant_id, @@ -411,8 +392,6 @@ def get_collections( ) if id: q = q.where(collections_t.id == ParameterValue(self.uuid_to_db(id))) - if topic: - q = q.where(collections_t.topic == ParameterValue(topic)) if name: q = q.where(collections_t.name == ParameterValue(name)) @@ -439,24 +418,24 @@ def get_collections( id = self.uuid_from_db(str(collection_id)) rows = list(collection_rows) name = str(rows[0][1]) - topic = str(rows[0][2]) - dimension = int(rows[0][3]) if rows[0][3] else None + dimension = int(rows[0][2]) if rows[0][2] else None metadata = self._metadata_from_rows(rows) collections.append( Collection( id=cast(UUID, id), - topic=topic, name=name, metadata=metadata, dimension=dimension, - tenant=str(rows[0][5]), - database=str(rows[0][4]), + tenant=str(rows[0][4]), + database=str(rows[0][3]), ) ) # apply limit and offset if limit is not None: - collections = collections[offset:offset+limit] + if offset is None: + offset = 0 + collections = collections[offset: offset + limit] else: collections = collections[offset:] @@ -494,7 +473,8 @@ def delete_collection( tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, ) -> None: - """Delete a topic and all associated segments from the SysDB""" + """Delete a collection and all associated segments from the SysDB. Deletes + the log stream for this collection as well.""" add_attributes_to_current_span( { "collection_id": str(id), @@ -519,18 +499,17 @@ def delete_collection( with self.tx() as cur: # no need for explicit del from metadata table because of ON DELETE CASCADE sql, params = get_sql(q, self.parameter_format()) - sql = sql + " RETURNING id, topic" + sql = sql + " RETURNING id" result = cur.execute(sql, params).fetchone() if not result: raise NotFoundError(f"Collection {id} not found") - self._producer.delete_topic(result[1]) + self._producer.delete_log(result[0]) @trace_method("SqlSysDB.update_segment", OpenTelemetryGranularity.ALL) @override def update_segment( self, id: UUID, - topic: OptionalArgument[Optional[str]] = Unspecified(), collection: OptionalArgument[Optional[UUID]] = Unspecified(), metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(), ) -> None: @@ -549,9 +528,6 @@ def update_segment( .where(segments_t.id == ParameterValue(self.uuid_to_db(id))) ) - if not topic == Unspecified(): - q = q.set(segments_t.topic, ParameterValue(topic)) - if not collection == Unspecified(): collection = cast(Optional[UUID], collection) q = q.set( @@ -589,7 +565,6 @@ def update_segment( def update_collection( self, id: UUID, - topic: OptionalArgument[Optional[str]] = Unspecified(), name: OptionalArgument[str] = Unspecified(), dimension: OptionalArgument[Optional[int]] = Unspecified(), metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(), @@ -608,9 +583,6 @@ def update_collection( .where(collections_t.id == ParameterValue(self.uuid_to_db(id))) ) - if not topic == Unspecified(): - q = q.set(collections_t.topic, ParameterValue(topic)) - if not name == Unspecified(): q = q.set(collections_t.name, ParameterValue(name)) diff --git a/chromadb/db/system.py b/chromadb/db/system.py index 15cbf5691c1..f971806ebe1 100644 --- a/chromadb/db/system.py +++ b/chromadb/db/system.py @@ -60,17 +60,15 @@ def get_segments( id: Optional[UUID] = None, type: Optional[str] = None, scope: Optional[SegmentScope] = None, - topic: Optional[str] = None, collection: Optional[UUID] = None, ) -> Sequence[Segment]: - """Find segments by id, type, scope, topic or collection.""" + """Find segments by id, type, scope or collection.""" pass @abstractmethod def update_segment( self, id: UUID, - topic: OptionalArgument[Optional[str]] = Unspecified(), collection: OptionalArgument[Optional[UUID]] = Unspecified(), metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(), ) -> None: @@ -91,8 +89,8 @@ def create_collection( database: str = DEFAULT_DATABASE, ) -> Tuple[Collection, bool]: """Create a new collection any associated resources - (Such as the necessary topics) in the SysDB. If get_or_create is True, the - collectionwill be created if one with the same name does not exist. + in the SysDB. If get_or_create is True, the + collection will be created if one with the same name does not exist. The metadata will be updated using the same protocol as update_collection. If get_or_create is False and the collection already exists, a error will be raised. @@ -105,7 +103,7 @@ def create_collection( def delete_collection( self, id: UUID, tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE ) -> None: - """Delete a collection, topic, all associated segments and any associate resources + """Delete a collection, all associated segments and any associate resources (log stream) from the SysDB and the system at large.""" pass @@ -113,21 +111,19 @@ def delete_collection( def get_collections( self, id: Optional[UUID] = None, - topic: Optional[str] = None, name: Optional[str] = None, tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE, limit: Optional[int] = None, offset: Optional[int] = None, ) -> Sequence[Collection]: - """Find collections by id, topic or name. If name is provided, tenant and database must also be provided.""" + """Find collections by id or name. If name is provided, tenant and database must also be provided.""" pass @abstractmethod def update_collection( self, id: UUID, - topic: OptionalArgument[str] = Unspecified(), name: OptionalArgument[str] = Unspecified(), dimension: OptionalArgument[Optional[int]] = Unspecified(), metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(), diff --git a/chromadb/ingest/__init__.py b/chromadb/ingest/__init__.py index 73f9cb065f2..32873369420 100644 --- a/chromadb/ingest/__init__.py +++ b/chromadb/ingest/__init__.py @@ -1,8 +1,8 @@ from abc import abstractmethod from typing import Callable, Optional, Sequence from chromadb.types import ( - SubmitEmbeddingRecord, - EmbeddingRecord, + OperationRecord, + LogRecord, SeqId, Vector, ScalarEncoding, @@ -38,25 +38,21 @@ class Producer(Component): """Interface for writing embeddings to an ingest stream""" @abstractmethod - def create_topic(self, topic_name: str) -> None: - pass - - @abstractmethod - def delete_topic(self, topic_name: str) -> None: + def delete_log(self, collection_id: UUID) -> None: pass @abstractmethod def submit_embedding( - self, topic_name: str, embedding: SubmitEmbeddingRecord + self, collection_id: UUID, embedding: OperationRecord ) -> SeqId: - """Add an embedding record to the given topic. Returns the SeqID of the record.""" + """Add an embedding record to the given collections log. Returns the SeqID of the record.""" pass @abstractmethod def submit_embeddings( - self, topic_name: str, embeddings: Sequence[SubmitEmbeddingRecord] + self, collection_id: UUID, embeddings: Sequence[OperationRecord] ) -> Sequence[SeqId]: - """Add a batch of embedding records to the given topic. Returns the SeqIDs of + """Add a batch of embedding records to the given collections log. Returns the SeqIDs of the records. The returned SeqIDs will be in the same order as the given SubmitEmbeddingRecords. However, it is not guaranteed that the SeqIDs will be processed in the same order as the given SubmitEmbeddingRecords. If the number @@ -71,7 +67,7 @@ def max_batch_size(self) -> int: pass -ConsumerCallbackFn = Callable[[Sequence[EmbeddingRecord]], None] +ConsumerCallbackFn = Callable[[Sequence[LogRecord]], None] class Consumer(Component): @@ -80,14 +76,14 @@ class Consumer(Component): @abstractmethod def subscribe( self, - topic_name: str, + collection_id: UUID, consume_fn: ConsumerCallbackFn, start: Optional[SeqId] = None, end: Optional[SeqId] = None, id: Optional[UUID] = None, ) -> UUID: """Register a function that will be called to recieve embeddings for a given - topic. The given function may be called any number of times, with any number of + collections log stream. The given function may be called any number of times, with any number of records, and may be called concurrently. Only records between start (exclusive) and end (inclusive) SeqIDs will be @@ -118,17 +114,3 @@ def min_seqid(self) -> SeqId: def max_seqid(self) -> SeqId: """Return the maximum possible SeqID in this implementation.""" pass - - -class CollectionAssignmentPolicy(Component): - """Interface for assigning collections to topics""" - - @abstractmethod - def assign_collection(self, collection_id: UUID) -> str: - """Return the topic that should be used for the given collection""" - pass - - @abstractmethod - def get_topics(self) -> Sequence[str]: - """Return the list of topics that this policy is currently using""" - pass diff --git a/chromadb/ingest/impl/pulsar.py b/chromadb/ingest/impl/pulsar.py deleted file mode 100644 index d84cadfa01e..00000000000 --- a/chromadb/ingest/impl/pulsar.py +++ /dev/null @@ -1,317 +0,0 @@ -from __future__ import annotations -from collections import defaultdict -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple -import uuid -from chromadb.config import Settings, System -from chromadb.ingest import Consumer, ConsumerCallbackFn, Producer -from overrides import overrides, EnforceOverrides -from uuid import UUID -from chromadb.ingest.impl.pulsar_admin import PulsarAdmin -from chromadb.ingest.impl.utils import create_pulsar_connection_str -from chromadb.proto.convert import from_proto_submit, to_proto_submit -import chromadb.proto.chroma_pb2 as proto -from chromadb.telemetry.opentelemetry import ( - OpenTelemetryClient, - OpenTelemetryGranularity, - trace_method, -) -from chromadb.types import SeqId, SubmitEmbeddingRecord -import pulsar -from concurrent.futures import wait, Future - -from chromadb.utils.messageid import int_to_pulsar, pulsar_to_int - - -class PulsarProducer(Producer, EnforceOverrides): - # TODO: ensure trace context propagates - _connection_str: str - _topic_to_producer: Dict[str, pulsar.Producer] - _opentelemetry_client: OpenTelemetryClient - _client: pulsar.Client - _admin: PulsarAdmin - _settings: Settings - - def __init__(self, system: System) -> None: - pulsar_host = system.settings.require("pulsar_broker_url") - pulsar_port = system.settings.require("pulsar_broker_port") - self._connection_str = create_pulsar_connection_str(pulsar_host, pulsar_port) - self._topic_to_producer = {} - self._settings = system.settings - self._admin = PulsarAdmin(system) - self._opentelemetry_client = system.require(OpenTelemetryClient) - super().__init__(system) - - @overrides - def start(self) -> None: - self._client = pulsar.Client(self._connection_str) - super().start() - - @overrides - def stop(self) -> None: - self._client.close() - super().stop() - - @overrides - def create_topic(self, topic_name: str) -> None: - self._admin.create_topic(topic_name) - - @overrides - def delete_topic(self, topic_name: str) -> None: - self._admin.delete_topic(topic_name) - - @trace_method("PulsarProducer.submit_embedding", OpenTelemetryGranularity.ALL) - @overrides - def submit_embedding( - self, topic_name: str, embedding: SubmitEmbeddingRecord - ) -> SeqId: - """Add an embedding record to the given topic. Returns the SeqID of the record.""" - producer = self._get_or_create_producer(topic_name) - proto_submit: proto.SubmitEmbeddingRecord = to_proto_submit(embedding) - # TODO: batch performance / async - msg_id: pulsar.MessageId = producer.send(proto_submit.SerializeToString()) - return pulsar_to_int(msg_id) - - @trace_method("PulsarProducer.submit_embeddings", OpenTelemetryGranularity.ALL) - @overrides - def submit_embeddings( - self, topic_name: str, embeddings: Sequence[SubmitEmbeddingRecord] - ) -> Sequence[SeqId]: - if not self._running: - raise RuntimeError("Component not running") - - if len(embeddings) == 0: - return [] - - if len(embeddings) > self.max_batch_size: - raise ValueError( - f""" - Cannot submit more than {self.max_batch_size:,} embeddings at once. - Please submit your embeddings in batches of size - {self.max_batch_size:,} or less. - """ - ) - - producer = self._get_or_create_producer(topic_name) - protos_to_submit = [to_proto_submit(embedding) for embedding in embeddings] - - def create_producer_callback( - future: Future[int], - ) -> Callable[[Any, pulsar.MessageId], None]: - def producer_callback(res: Any, msg_id: pulsar.MessageId) -> None: - if msg_id: - future.set_result(pulsar_to_int(msg_id)) - else: - future.set_exception( - Exception( - "Unknown error while submitting embedding in producer_callback" - ) - ) - - return producer_callback - - futures = [] - for proto_to_submit in protos_to_submit: - future: Future[int] = Future() - producer.send_async( - proto_to_submit.SerializeToString(), - callback=create_producer_callback(future), - ) - futures.append(future) - - wait(futures) - - results: List[SeqId] = [] - for future in futures: - exception = future.exception() - if exception is not None: - raise exception - results.append(future.result()) - - return results - - @property - @overrides - def max_batch_size(self) -> int: - # For now, we use 1,000 - # TODO: tune this to a reasonable value by default - return 1000 - - def _get_or_create_producer(self, topic_name: str) -> pulsar.Producer: - if topic_name not in self._topic_to_producer: - producer = self._client.create_producer(topic_name) - self._topic_to_producer[topic_name] = producer - return self._topic_to_producer[topic_name] - - @overrides - def reset_state(self) -> None: - if not self._settings.require("allow_reset"): - raise ValueError( - "Resetting the database is not allowed. Set `allow_reset` to true in the config in tests or other non-production environments where reset should be permitted." - ) - for topic_name in self._topic_to_producer: - self._admin.delete_topic(topic_name) - self._topic_to_producer = {} - super().reset_state() - - -class PulsarConsumer(Consumer, EnforceOverrides): - class PulsarSubscription: - id: UUID - topic_name: str - start: int - end: int - callback: ConsumerCallbackFn - consumer: pulsar.Consumer - - def __init__( - self, - id: UUID, - topic_name: str, - start: int, - end: int, - callback: ConsumerCallbackFn, - consumer: pulsar.Consumer, - ): - self.id = id - self.topic_name = topic_name - self.start = start - self.end = end - self.callback = callback - self.consumer = consumer - - _connection_str: str - _client: pulsar.Client - _opentelemetry_client: OpenTelemetryClient - _subscriptions: Dict[str, Set[PulsarSubscription]] - _settings: Settings - - def __init__(self, system: System) -> None: - pulsar_host = system.settings.require("pulsar_broker_url") - pulsar_port = system.settings.require("pulsar_broker_port") - self._connection_str = create_pulsar_connection_str(pulsar_host, pulsar_port) - self._subscriptions = defaultdict(set) - self._settings = system.settings - self._opentelemetry_client = system.require(OpenTelemetryClient) - super().__init__(system) - - @overrides - def start(self) -> None: - self._client = pulsar.Client(self._connection_str) - super().start() - - @overrides - def stop(self) -> None: - self._client.close() - super().stop() - - @trace_method("PulsarConsumer.subscribe", OpenTelemetryGranularity.ALL) - @overrides - def subscribe( - self, - topic_name: str, - consume_fn: ConsumerCallbackFn, - start: Optional[SeqId] = None, - end: Optional[SeqId] = None, - id: Optional[UUID] = None, - ) -> UUID: - """Register a function that will be called to recieve embeddings for a given - topic. The given function may be called any number of times, with any number of - records, and may be called concurrently. - - Only records between start (exclusive) and end (inclusive) SeqIDs will be - returned. If start is None, the first record returned will be the next record - generated, not including those generated before creating the subscription. If - end is None, the consumer will consume indefinitely, otherwise it will - automatically be unsubscribed when the end SeqID is reached. - - If the function throws an exception, the function may be called again with the - same or different records. - - Takes an optional UUID as a unique subscription ID. If no ID is provided, a new - ID will be generated and returned.""" - if not self._running: - raise RuntimeError("Consumer must be started before subscribing") - - subscription_id = ( - id or uuid.uuid4() - ) # TODO: this should really be created by the coordinator and stored in sysdb - - start, end = self._validate_range(start, end) - - def wrap_callback(consumer: pulsar.Consumer, message: pulsar.Message) -> None: - msg_data = message.data() - msg_id = pulsar_to_int(message.message_id()) - submit_embedding_record = proto.SubmitEmbeddingRecord() - proto.SubmitEmbeddingRecord.ParseFromString( - submit_embedding_record, msg_data - ) - embedding_record = from_proto_submit(submit_embedding_record, msg_id) - consume_fn([embedding_record]) - consumer.acknowledge(message) - if msg_id == end: - self.unsubscribe(subscription_id) - - consumer = self._client.subscribe( - topic_name, - subscription_id.hex, - message_listener=wrap_callback, - ) - - subscription = self.PulsarSubscription( - subscription_id, topic_name, start, end, consume_fn, consumer - ) - self._subscriptions[topic_name].add(subscription) - - # NOTE: For some reason the seek() method expects a shadowed MessageId type - # which resides in _msg_id. - consumer.seek(int_to_pulsar(start)._msg_id) - - return subscription_id - - def _validate_range( - self, start: Optional[SeqId], end: Optional[SeqId] - ) -> Tuple[int, int]: - """Validate and normalize the start and end SeqIDs for a subscription using this - impl.""" - start = start or pulsar_to_int(pulsar.MessageId.latest) - end = end or self.max_seqid() - if not isinstance(start, int) or not isinstance(end, int): - raise TypeError("SeqIDs must be integers") - if start >= end: - raise ValueError(f"Invalid SeqID range: {start} to {end}") - return start, end - - @overrides - def unsubscribe(self, subscription_id: UUID) -> None: - """Unregister a subscription. The consume function will no longer be invoked, - and resources associated with the subscription will be released.""" - for topic_name, subscriptions in self._subscriptions.items(): - for subscription in subscriptions: - if subscription.id == subscription_id: - subscription.consumer.close() - subscriptions.remove(subscription) - if len(subscriptions) == 0: - del self._subscriptions[topic_name] - return - - @overrides - def min_seqid(self) -> SeqId: - """Return the minimum possible SeqID in this implementation.""" - return pulsar_to_int(pulsar.MessageId.earliest) - - @overrides - def max_seqid(self) -> SeqId: - """Return the maximum possible SeqID in this implementation.""" - return 2**192 - 1 - - @overrides - def reset_state(self) -> None: - if not self._settings.require("allow_reset"): - raise ValueError( - "Resetting the database is not allowed. Set `allow_reset` to true in the config in tests or other non-production environments where reset should be permitted." - ) - for topic_name, subscriptions in self._subscriptions.items(): - for subscription in subscriptions: - subscription.consumer.close() - self._subscriptions = defaultdict(set) - super().reset_state() diff --git a/chromadb/ingest/impl/pulsar_admin.py b/chromadb/ingest/impl/pulsar_admin.py deleted file mode 100644 index e031e4a238b..00000000000 --- a/chromadb/ingest/impl/pulsar_admin.py +++ /dev/null @@ -1,81 +0,0 @@ -# A thin wrapper around the pulsar admin api -import requests -from chromadb.config import System -from chromadb.ingest.impl.utils import parse_topic_name - - -class PulsarAdmin: - """A thin wrapper around the pulsar admin api, only used for interim development towards distributed chroma. - This functionality will be moved to the chroma coordinator.""" - - _connection_str: str - - def __init__(self, system: System): - pulsar_host = system.settings.require("pulsar_broker_url") - pulsar_port = system.settings.require("pulsar_admin_port") - self._connection_str = f"http://{pulsar_host}:{pulsar_port}" - - # Create the default tenant and namespace - # This is a temporary workaround until we have a proper tenant/namespace management system - self.create_tenant("default") - self.create_namespace("default", "default") - - def create_tenant(self, tenant: str) -> None: - """Make a PUT request to the admin api to create the tenant""" - - path = f"/admin/v2/tenants/{tenant}" - url = self._connection_str + path - response = requests.put( - url, json={"allowedClusters": ["standalone"], "adminRoles": []} - ) # TODO: how to manage clusters? - - if response.status_code != 204 and response.status_code != 409: - raise RuntimeError(f"Failed to create tenant {tenant}") - - def create_namespace(self, tenant: str, namespace: str) -> None: - """Make a PUT request to the admin api to create the namespace""" - - path = f"/admin/v2/namespaces/{tenant}/{namespace}" - url = self._connection_str + path - response = requests.put(url) - - if response.status_code != 204 and response.status_code != 409: - raise RuntimeError(f"Failed to create namespace {namespace}") - - def create_topic(self, topic: str) -> None: - # TODO: support non-persistent topics? - tenant, namespace, topic_name = parse_topic_name(topic) - - if tenant != "default": - raise ValueError(f"Only the default tenant is supported, got {tenant}") - if namespace != "default": - raise ValueError( - f"Only the default namespace is supported, got {namespace}" - ) - - # Make a PUT request to the admin api to create the topic - path = f"/admin/v2/persistent/{tenant}/{namespace}/{topic_name}" - url = self._connection_str + path - response = requests.put(url) - - if response.status_code != 204 and response.status_code != 409: - raise RuntimeError(f"Failed to create topic {topic_name}") - - def delete_topic(self, topic: str) -> None: - tenant, namespace, topic_name = parse_topic_name(topic) - - if tenant != "default": - raise ValueError(f"Only the default tenant is supported, got {tenant}") - if namespace != "default": - raise ValueError( - f"Only the default namespace is supported, got {namespace}" - ) - - # Make a PUT request to the admin api to delete the topic - path = f"/admin/v2/persistent/{tenant}/{namespace}/{topic_name}" - # Force delete the topic - path += "?force=true" - url = self._connection_str + path - response = requests.delete(url) - if response.status_code != 204 and response.status_code != 409: - raise RuntimeError(f"Failed to delete topic {topic_name}") diff --git a/chromadb/ingest/impl/simple_policy.py b/chromadb/ingest/impl/simple_policy.py deleted file mode 100644 index f8068ee2046..00000000000 --- a/chromadb/ingest/impl/simple_policy.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import Sequence -from uuid import UUID -from overrides import overrides -from chromadb.config import System -from chromadb.ingest import CollectionAssignmentPolicy -from chromadb.ingest.impl.utils import create_topic_name - - -class SimpleAssignmentPolicy(CollectionAssignmentPolicy): - """Simple assignment policy that assigns a 1 collection to 1 topic based on the - id of the collection.""" - - _tenant_id: str - _topic_ns: str - - def __init__(self, system: System): - self._tenant_id = system.settings.tenant_id - self._topic_ns = system.settings.topic_namespace - super().__init__(system) - - def _topic(self, collection_id: UUID) -> str: - return create_topic_name(self._tenant_id, self._topic_ns, str(collection_id)) - - @overrides - def assign_collection(self, collection_id: UUID) -> str: - return self._topic(collection_id) - - @overrides - def get_topics(self) -> Sequence[str]: - raise NotImplementedError( - "SimpleAssignmentPolicy does not support get_topics, each collection has its own topic" - ) - - -class RendezvousHashingAssignmentPolicy(CollectionAssignmentPolicy): - """The rendezvous hashing assignment policy assigns a collection to a topic based on the - rendezvous hashing algorithm. This is not actually used in the python sysdb. It is only used in the - go sysdb. However, it is useful here in order to provide a way to get the topic list used for the whole system. - """ - - _tenant_id: str - _topic_ns: str - - def __init__(self, system: System): - self._tenant_id = system.settings.tenant_id - self._topic_ns = system.settings.topic_namespace - super().__init__(system) - - @overrides - def assign_collection(self, collection_id: UUID) -> str: - raise NotImplementedError( - "RendezvousHashingAssignmentPolicy is not implemented" - ) - - @overrides - def get_topics(self) -> Sequence[str]: - # Mirrors go/coordinator/internal/coordinator/assignment_policy.go - return [ - f"persistent://{self._tenant_id}/{self._topic_ns}/chroma_log_{i}" - for i in range(16) - ] diff --git a/chromadb/ingest/impl/utils.py b/chromadb/ingest/impl/utils.py index 144384d75db..b28a7d8a112 100644 --- a/chromadb/ingest/impl/utils.py +++ b/chromadb/ingest/impl/utils.py @@ -1,5 +1,6 @@ import re from typing import Tuple +from uuid import UUID topic_regex = r"persistent:\/\/(?P.+)\/(?P.+)\/(?P.+)" @@ -12,9 +13,5 @@ def parse_topic_name(topic_name: str) -> Tuple[str, str, str]: return match.group("tenant"), match.group("namespace"), match.group("topic") -def create_pulsar_connection_str(host: str, port: str) -> str: - return f"pulsar://{host}:{port}" - - -def create_topic_name(tenant: str, namespace: str, topic: str) -> str: - return f"persistent://{tenant}/{namespace}/{topic}" +def create_topic_name(tenant: str, namespace: str, collection_id: UUID) -> str: + return f"persistent://{tenant}/{namespace}/{str(collection_id)}" diff --git a/chromadb/logservice/logservice.py b/chromadb/logservice/logservice.py new file mode 100644 index 00000000000..4024ec9b277 --- /dev/null +++ b/chromadb/logservice/logservice.py @@ -0,0 +1,160 @@ +import sys + +import grpc + +from chromadb.ingest import ( + Producer, + Consumer, + ConsumerCallbackFn, +) +from chromadb.proto.convert import to_proto_submit +from chromadb.proto.logservice_pb2 import PushLogsRequest, PullLogsRequest, LogRecord +from chromadb.proto.logservice_pb2_grpc import LogServiceStub +from chromadb.telemetry.opentelemetry.grpc import OtelInterceptor +from chromadb.types import ( + OperationRecord, + SeqId, +) +from chromadb.config import System +from chromadb.telemetry.opentelemetry import ( + OpenTelemetryClient, + OpenTelemetryGranularity, + trace_method, +) +from overrides import override +from typing import Sequence, Optional, cast +from uuid import UUID +import logging + +logger = logging.getLogger(__name__) + + +class LogService(Producer, Consumer): + """ + Distributed Chroma Log Service + """ + + _log_service_stub: LogServiceStub + _channel: grpc.Channel + _log_service_url: str + _log_service_port: int + + def __init__(self, system: System): + self._log_service_url = system.settings.require("chroma_logservice_host") + self._log_service_port = system.settings.require("chroma_logservice_port") + self._opentelemetry_client = system.require(OpenTelemetryClient) + super().__init__(system) + + @trace_method("LogService.start", OpenTelemetryGranularity.ALL) + @override + def start(self) -> None: + self._channel = grpc.insecure_channel( + f"{self._log_service_url}:{self._log_service_port}" + ) + interceptors = [OtelInterceptor()] + self._channel = grpc.intercept_channel(self._channel, *interceptors) + self._log_service_stub = LogServiceStub(self._channel) # type: ignore + super().start() + + @trace_method("LogService.stop", OpenTelemetryGranularity.ALL) + @override + def stop(self) -> None: + self._channel.close() + super().stop() + + @trace_method("LogService.reset_state", OpenTelemetryGranularity.ALL) + @override + def reset_state(self) -> None: + super().reset_state() + + @trace_method("LogService.delete_log", OpenTelemetryGranularity.ALL) + @override + def delete_log(self, collection_id: UUID) -> None: + raise NotImplementedError("Not implemented") + + @trace_method("LogService.submit_embedding", OpenTelemetryGranularity.ALL) + @override + def submit_embedding( + self, collection_id: UUID, embedding: OperationRecord + ) -> SeqId: + if not self._running: + raise RuntimeError("Component not running") + + return self.submit_embeddings(collection_id, [embedding])[0] # type: ignore + + @trace_method("LogService.submit_embeddings", OpenTelemetryGranularity.ALL) + @override + def submit_embeddings( + self, collection_id: UUID, embeddings: Sequence[OperationRecord] + ) -> Sequence[SeqId]: + logger.info( + f"Submitting {len(embeddings)} embeddings to log for collection {collection_id}" + ) + + if not self._running: + raise RuntimeError("Component not running") + + if len(embeddings) == 0: + return [] + + # push records to the log service + counts = [] + protos_to_submit = [to_proto_submit(record) for record in embeddings] + counts.append( + self.push_logs( + collection_id, + cast(Sequence[OperationRecord], protos_to_submit), + ) + ) + + # This returns counts, which is completely incorrect + # TODO: Fix this + return counts + + @trace_method("LogService.subscribe", OpenTelemetryGranularity.ALL) + @override + def subscribe( + self, + collection_id: UUID, + consume_fn: ConsumerCallbackFn, + start: Optional[SeqId] = None, + end: Optional[SeqId] = None, + id: Optional[UUID] = None, + ) -> UUID: + logger.info(f"Subscribing to log for {collection_id}, noop for logservice") + return UUID(int=0) + + @trace_method("LogService.unsubscribe", OpenTelemetryGranularity.ALL) + @override + def unsubscribe(self, subscription_id: UUID) -> None: + logger.info(f"Unsubscribing from {subscription_id}, noop for logservice") + + @override + def min_seqid(self) -> SeqId: + return 0 + + @override + def max_seqid(self) -> SeqId: + return sys.maxsize + + @property + @override + def max_batch_size(self) -> int: + return sys.maxsize + + def push_logs(self, collection_id: UUID, records: Sequence[OperationRecord]) -> int: + request = PushLogsRequest(collection_id=str(collection_id), records=records) + response = self._log_service_stub.PushLogs(request) + return response.record_count # type: ignore + + def pull_logs( + self, collection_id: UUID, start_offset: int, batch_size: int + ) -> Sequence[LogRecord]: + request = PullLogsRequest( + collection_id=str(collection_id), + start_from_offset=start_offset, + batch_size=batch_size, + end_timestamp=-1, + ) + response = self._log_service_stub.PullLogs(request) + return response.records # type: ignore diff --git a/chromadb/migrations/sysdb/00005-remove-topic.sqlite.sql b/chromadb/migrations/sysdb/00005-remove-topic.sqlite.sql new file mode 100644 index 00000000000..3ed0e028423 --- /dev/null +++ b/chromadb/migrations/sysdb/00005-remove-topic.sqlite.sql @@ -0,0 +1,4 @@ +-- Remove the topic column from the Collections and Segments tables + +ALTER TABLE collections DROP COLUMN topic; +ALTER TABLE segments DROP COLUMN topic; diff --git a/chromadb/proto/chroma_pb2.py b/chromadb/proto/chroma_pb2.py index 84a3ba9b13d..333b2e70a66 100644 --- a/chromadb/proto/chroma_pb2.py +++ b/chromadb/proto/chroma_pb2.py @@ -13,62 +13,62 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x63hromadb/proto/chroma.proto\x12\x06\x63hroma\"&\n\x06Status\x12\x0e\n\x06reason\x18\x01 \x01(\t\x12\x0c\n\x04\x63ode\x18\x02 \x01(\x05\"0\n\x0e\x43hromaResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"U\n\x06Vector\x12\x11\n\tdimension\x18\x01 \x01(\x05\x12\x0e\n\x06vector\x18\x02 \x01(\x0c\x12(\n\x08\x65ncoding\x18\x03 \x01(\x0e\x32\x16.chroma.ScalarEncoding\"\xca\x01\n\x07Segment\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12#\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScope\x12\x12\n\x05topic\x18\x04 \x01(\tH\x00\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x01\x88\x01\x01\x12-\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x02\x88\x01\x01\x42\x08\n\x06_topicB\r\n\x0b_collectionB\x0b\n\t_metadata\"\xb9\x01\n\nCollection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\r\n\x05topic\x18\x03 \x01(\t\x12-\n\x08metadata\x18\x04 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x05 \x01(\x05H\x01\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\tB\x0b\n\t_metadataB\x0c\n\n_dimension\"4\n\x08\x44\x61tabase\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"\x16\n\x06Tenant\x12\x0c\n\x04name\x18\x01 \x01(\t\"b\n\x13UpdateMetadataValue\x12\x16\n\x0cstring_value\x18\x01 \x01(\tH\x00\x12\x13\n\tint_value\x18\x02 \x01(\x03H\x00\x12\x15\n\x0b\x66loat_value\x18\x03 \x01(\x01H\x00\x42\x07\n\x05value\"\x96\x01\n\x0eUpdateMetadata\x12\x36\n\x08metadata\x18\x01 \x03(\x0b\x32$.chroma.UpdateMetadata.MetadataEntry\x1aL\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12*\n\x05value\x18\x02 \x01(\x0b\x32\x1b.chroma.UpdateMetadataValue:\x02\x38\x01\"\xcc\x01\n\x15SubmitEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12#\n\x06vector\x18\x02 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x01\x88\x01\x01\x12$\n\toperation\x18\x04 \x01(\x0e\x32\x11.chroma.Operation\x12\x15\n\rcollection_id\x18\x05 \x01(\tB\t\n\x07_vectorB\x0b\n\t_metadata\"S\n\x15VectorEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06seq_id\x18\x02 \x01(\x0c\x12\x1e\n\x06vector\x18\x03 \x01(\x0b\x32\x0e.chroma.Vector\"q\n\x11VectorQueryResult\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06seq_id\x18\x02 \x01(\x0c\x12\x10\n\x08\x64istance\x18\x03 \x01(\x01\x12#\n\x06vector\x18\x04 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x42\t\n\x07_vector\"@\n\x12VectorQueryResults\x12*\n\x07results\x18\x01 \x03(\x0b\x32\x19.chroma.VectorQueryResult\"(\n\x15SegmentServerResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\"4\n\x11GetVectorsRequest\x12\x0b\n\x03ids\x18\x01 \x03(\t\x12\x12\n\nsegment_id\x18\x02 \x01(\t\"D\n\x12GetVectorsResponse\x12.\n\x07records\x18\x01 \x03(\x0b\x32\x1d.chroma.VectorEmbeddingRecord\"\x86\x01\n\x13QueryVectorsRequest\x12\x1f\n\x07vectors\x18\x01 \x03(\x0b\x32\x0e.chroma.Vector\x12\t\n\x01k\x18\x02 \x01(\x05\x12\x13\n\x0b\x61llowed_ids\x18\x03 \x03(\t\x12\x1a\n\x12include_embeddings\x18\x04 \x01(\x08\x12\x12\n\nsegment_id\x18\x05 \x01(\t\"C\n\x14QueryVectorsResponse\x12+\n\x07results\x18\x01 \x03(\x0b\x32\x1a.chroma.VectorQueryResults*8\n\tOperation\x12\x07\n\x03\x41\x44\x44\x10\x00\x12\n\n\x06UPDATE\x10\x01\x12\n\n\x06UPSERT\x10\x02\x12\n\n\x06\x44\x45LETE\x10\x03*(\n\x0eScalarEncoding\x12\x0b\n\x07\x46LOAT32\x10\x00\x12\t\n\x05INT32\x10\x01*(\n\x0cSegmentScope\x12\n\n\x06VECTOR\x10\x00\x12\x0c\n\x08METADATA\x10\x01\x32\x94\x01\n\rSegmentServer\x12?\n\x0bLoadSegment\x12\x0f.chroma.Segment\x1a\x1d.chroma.SegmentServerResponse\"\x00\x12\x42\n\x0eReleaseSegment\x12\x0f.chroma.Segment\x1a\x1d.chroma.SegmentServerResponse\"\x00\x32\xa2\x01\n\x0cVectorReader\x12\x45\n\nGetVectors\x12\x19.chroma.GetVectorsRequest\x1a\x1a.chroma.GetVectorsResponse\"\x00\x12K\n\x0cQueryVectors\x12\x1b.chroma.QueryVectorsRequest\x1a\x1c.chroma.QueryVectorsResponse\"\x00\x42\x43ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpbb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x63hromadb/proto/chroma.proto\x12\x06\x63hroma\"&\n\x06Status\x12\x0e\n\x06reason\x18\x01 \x01(\t\x12\x0c\n\x04\x63ode\x18\x02 \x01(\x05\"U\n\x06Vector\x12\x11\n\tdimension\x18\x01 \x01(\x05\x12\x0e\n\x06vector\x18\x02 \x01(\x0c\x12(\n\x08\x65ncoding\x18\x03 \x01(\x0e\x32\x16.chroma.ScalarEncoding\"\x1a\n\tFilePaths\x12\r\n\x05paths\x18\x01 \x03(\t\"\xa5\x02\n\x07Segment\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12#\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScope\x12\x17\n\ncollection\x18\x05 \x01(\tH\x00\x88\x01\x01\x12-\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x01\x88\x01\x01\x12\x32\n\nfile_paths\x18\x07 \x03(\x0b\x32\x1e.chroma.Segment.FilePathsEntry\x1a\x43\n\x0e\x46ilePathsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12 \n\x05value\x18\x02 \x01(\x0b\x32\x11.chroma.FilePaths:\x02\x38\x01\x42\r\n\x0b_collectionB\x0b\n\t_metadata\"\xd1\x01\n\nCollection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12-\n\x08metadata\x18\x04 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x05 \x01(\x05H\x01\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\t\x12\x14\n\x0clog_position\x18\x08 \x01(\x03\x12\x0f\n\x07version\x18\t \x01(\x05\x42\x0b\n\t_metadataB\x0c\n\n_dimension\"4\n\x08\x44\x61tabase\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"\x16\n\x06Tenant\x12\x0c\n\x04name\x18\x01 \x01(\t\"b\n\x13UpdateMetadataValue\x12\x16\n\x0cstring_value\x18\x01 \x01(\tH\x00\x12\x13\n\tint_value\x18\x02 \x01(\x03H\x00\x12\x15\n\x0b\x66loat_value\x18\x03 \x01(\x01H\x00\x42\x07\n\x05value\"\x96\x01\n\x0eUpdateMetadata\x12\x36\n\x08metadata\x18\x01 \x03(\x0b\x32$.chroma.UpdateMetadata.MetadataEntry\x1aL\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12*\n\x05value\x18\x02 \x01(\x0b\x32\x1b.chroma.UpdateMetadataValue:\x02\x38\x01\"\xaf\x01\n\x0fOperationRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12#\n\x06vector\x18\x02 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x01\x88\x01\x01\x12$\n\toperation\x18\x04 \x01(\x0e\x32\x11.chroma.OperationB\t\n\x07_vectorB\x0b\n\t_metadata\"C\n\x15VectorEmbeddingRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x1e\n\x06vector\x18\x03 \x01(\x0b\x32\x0e.chroma.Vector\"a\n\x11VectorQueryResult\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x64istance\x18\x03 \x01(\x02\x12#\n\x06vector\x18\x04 \x01(\x0b\x32\x0e.chroma.VectorH\x00\x88\x01\x01\x42\t\n\x07_vector\"@\n\x12VectorQueryResults\x12*\n\x07results\x18\x01 \x03(\x0b\x32\x19.chroma.VectorQueryResult\"4\n\x11GetVectorsRequest\x12\x0b\n\x03ids\x18\x01 \x03(\t\x12\x12\n\nsegment_id\x18\x02 \x01(\t\"D\n\x12GetVectorsResponse\x12.\n\x07records\x18\x01 \x03(\x0b\x32\x1d.chroma.VectorEmbeddingRecord\"\x86\x01\n\x13QueryVectorsRequest\x12\x1f\n\x07vectors\x18\x01 \x03(\x0b\x32\x0e.chroma.Vector\x12\t\n\x01k\x18\x02 \x01(\x05\x12\x13\n\x0b\x61llowed_ids\x18\x03 \x03(\t\x12\x1a\n\x12include_embeddings\x18\x04 \x01(\x08\x12\x12\n\nsegment_id\x18\x05 \x01(\t\"C\n\x14QueryVectorsResponse\x12+\n\x07results\x18\x01 \x03(\x0b\x32\x1a.chroma.VectorQueryResults*8\n\tOperation\x12\x07\n\x03\x41\x44\x44\x10\x00\x12\n\n\x06UPDATE\x10\x01\x12\n\n\x06UPSERT\x10\x02\x12\n\n\x06\x44\x45LETE\x10\x03*(\n\x0eScalarEncoding\x12\x0b\n\x07\x46LOAT32\x10\x00\x12\t\n\x05INT32\x10\x01*(\n\x0cSegmentScope\x12\n\n\x06VECTOR\x10\x00\x12\x0c\n\x08METADATA\x10\x01\x32\xa2\x01\n\x0cVectorReader\x12\x45\n\nGetVectors\x12\x19.chroma.GetVectorsRequest\x1a\x1a.chroma.GetVectorsResponse\"\x00\x12K\n\x0cQueryVectors\x12\x1b.chroma.QueryVectorsRequest\x1a\x1c.chroma.QueryVectorsResponse\"\x00\x42:Z8github.com/chroma-core/chroma/go/pkg/proto/coordinatorpbb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'chromadb.proto.chroma_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: DESCRIPTOR._options = None - DESCRIPTOR._serialized_options = b'ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpb' + DESCRIPTOR._serialized_options = b'Z8github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb' + _SEGMENT_FILEPATHSENTRY._options = None + _SEGMENT_FILEPATHSENTRY._serialized_options = b'8\001' _UPDATEMETADATA_METADATAENTRY._options = None _UPDATEMETADATA_METADATAENTRY._serialized_options = b'8\001' - _globals['_OPERATION']._serialized_start=1785 - _globals['_OPERATION']._serialized_end=1841 - _globals['_SCALARENCODING']._serialized_start=1843 - _globals['_SCALARENCODING']._serialized_end=1883 - _globals['_SEGMENTSCOPE']._serialized_start=1885 - _globals['_SEGMENTSCOPE']._serialized_end=1925 + _globals['_OPERATION']._serialized_start=1775 + _globals['_OPERATION']._serialized_end=1831 + _globals['_SCALARENCODING']._serialized_start=1833 + _globals['_SCALARENCODING']._serialized_end=1873 + _globals['_SEGMENTSCOPE']._serialized_start=1875 + _globals['_SEGMENTSCOPE']._serialized_end=1915 _globals['_STATUS']._serialized_start=39 _globals['_STATUS']._serialized_end=77 - _globals['_CHROMARESPONSE']._serialized_start=79 - _globals['_CHROMARESPONSE']._serialized_end=127 - _globals['_VECTOR']._serialized_start=129 - _globals['_VECTOR']._serialized_end=214 - _globals['_SEGMENT']._serialized_start=217 - _globals['_SEGMENT']._serialized_end=419 - _globals['_COLLECTION']._serialized_start=422 - _globals['_COLLECTION']._serialized_end=607 - _globals['_DATABASE']._serialized_start=609 - _globals['_DATABASE']._serialized_end=661 - _globals['_TENANT']._serialized_start=663 - _globals['_TENANT']._serialized_end=685 - _globals['_UPDATEMETADATAVALUE']._serialized_start=687 - _globals['_UPDATEMETADATAVALUE']._serialized_end=785 - _globals['_UPDATEMETADATA']._serialized_start=788 - _globals['_UPDATEMETADATA']._serialized_end=938 - _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_start=862 - _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_end=938 - _globals['_SUBMITEMBEDDINGRECORD']._serialized_start=941 - _globals['_SUBMITEMBEDDINGRECORD']._serialized_end=1145 - _globals['_VECTOREMBEDDINGRECORD']._serialized_start=1147 - _globals['_VECTOREMBEDDINGRECORD']._serialized_end=1230 - _globals['_VECTORQUERYRESULT']._serialized_start=1232 - _globals['_VECTORQUERYRESULT']._serialized_end=1345 - _globals['_VECTORQUERYRESULTS']._serialized_start=1347 - _globals['_VECTORQUERYRESULTS']._serialized_end=1411 - _globals['_SEGMENTSERVERRESPONSE']._serialized_start=1413 - _globals['_SEGMENTSERVERRESPONSE']._serialized_end=1453 - _globals['_GETVECTORSREQUEST']._serialized_start=1455 - _globals['_GETVECTORSREQUEST']._serialized_end=1507 - _globals['_GETVECTORSRESPONSE']._serialized_start=1509 - _globals['_GETVECTORSRESPONSE']._serialized_end=1577 - _globals['_QUERYVECTORSREQUEST']._serialized_start=1580 - _globals['_QUERYVECTORSREQUEST']._serialized_end=1714 - _globals['_QUERYVECTORSRESPONSE']._serialized_start=1716 - _globals['_QUERYVECTORSRESPONSE']._serialized_end=1783 - _globals['_SEGMENTSERVER']._serialized_start=1928 - _globals['_SEGMENTSERVER']._serialized_end=2076 - _globals['_VECTORREADER']._serialized_start=2079 - _globals['_VECTORREADER']._serialized_end=2241 + _globals['_VECTOR']._serialized_start=79 + _globals['_VECTOR']._serialized_end=164 + _globals['_FILEPATHS']._serialized_start=166 + _globals['_FILEPATHS']._serialized_end=192 + _globals['_SEGMENT']._serialized_start=195 + _globals['_SEGMENT']._serialized_end=488 + _globals['_SEGMENT_FILEPATHSENTRY']._serialized_start=393 + _globals['_SEGMENT_FILEPATHSENTRY']._serialized_end=460 + _globals['_COLLECTION']._serialized_start=491 + _globals['_COLLECTION']._serialized_end=700 + _globals['_DATABASE']._serialized_start=702 + _globals['_DATABASE']._serialized_end=754 + _globals['_TENANT']._serialized_start=756 + _globals['_TENANT']._serialized_end=778 + _globals['_UPDATEMETADATAVALUE']._serialized_start=780 + _globals['_UPDATEMETADATAVALUE']._serialized_end=878 + _globals['_UPDATEMETADATA']._serialized_start=881 + _globals['_UPDATEMETADATA']._serialized_end=1031 + _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_start=955 + _globals['_UPDATEMETADATA_METADATAENTRY']._serialized_end=1031 + _globals['_OPERATIONRECORD']._serialized_start=1034 + _globals['_OPERATIONRECORD']._serialized_end=1209 + _globals['_VECTOREMBEDDINGRECORD']._serialized_start=1211 + _globals['_VECTOREMBEDDINGRECORD']._serialized_end=1278 + _globals['_VECTORQUERYRESULT']._serialized_start=1280 + _globals['_VECTORQUERYRESULT']._serialized_end=1377 + _globals['_VECTORQUERYRESULTS']._serialized_start=1379 + _globals['_VECTORQUERYRESULTS']._serialized_end=1443 + _globals['_GETVECTORSREQUEST']._serialized_start=1445 + _globals['_GETVECTORSREQUEST']._serialized_end=1497 + _globals['_GETVECTORSRESPONSE']._serialized_start=1499 + _globals['_GETVECTORSRESPONSE']._serialized_end=1567 + _globals['_QUERYVECTORSREQUEST']._serialized_start=1570 + _globals['_QUERYVECTORSREQUEST']._serialized_end=1704 + _globals['_QUERYVECTORSRESPONSE']._serialized_start=1706 + _globals['_QUERYVECTORSRESPONSE']._serialized_end=1773 + _globals['_VECTORREADER']._serialized_start=1918 + _globals['_VECTORREADER']._serialized_end=2080 # @@protoc_insertion_point(module_scope) diff --git a/chromadb/proto/chroma_pb2.pyi b/chromadb/proto/chroma_pb2.pyi index 026bfac8821..997dc12c3f4 100644 --- a/chromadb/proto/chroma_pb2.pyi +++ b/chromadb/proto/chroma_pb2.pyi @@ -39,12 +39,6 @@ class Status(_message.Message): code: int def __init__(self, reason: _Optional[str] = ..., code: _Optional[int] = ...) -> None: ... -class ChromaResponse(_message.Message): - __slots__ = ["status"] - STATUS_FIELD_NUMBER: _ClassVar[int] - status: Status - def __init__(self, status: _Optional[_Union[Status, _Mapping]] = ...) -> None: ... - class Vector(_message.Message): __slots__ = ["dimension", "vector", "encoding"] DIMENSION_FIELD_NUMBER: _ClassVar[int] @@ -55,39 +49,54 @@ class Vector(_message.Message): encoding: ScalarEncoding def __init__(self, dimension: _Optional[int] = ..., vector: _Optional[bytes] = ..., encoding: _Optional[_Union[ScalarEncoding, str]] = ...) -> None: ... +class FilePaths(_message.Message): + __slots__ = ["paths"] + PATHS_FIELD_NUMBER: _ClassVar[int] + paths: _containers.RepeatedScalarFieldContainer[str] + def __init__(self, paths: _Optional[_Iterable[str]] = ...) -> None: ... + class Segment(_message.Message): - __slots__ = ["id", "type", "scope", "topic", "collection", "metadata"] + __slots__ = ["id", "type", "scope", "collection", "metadata", "file_paths"] + class FilePathsEntry(_message.Message): + __slots__ = ["key", "value"] + KEY_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] + key: str + value: FilePaths + def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[FilePaths, _Mapping]] = ...) -> None: ... ID_FIELD_NUMBER: _ClassVar[int] TYPE_FIELD_NUMBER: _ClassVar[int] SCOPE_FIELD_NUMBER: _ClassVar[int] - TOPIC_FIELD_NUMBER: _ClassVar[int] COLLECTION_FIELD_NUMBER: _ClassVar[int] METADATA_FIELD_NUMBER: _ClassVar[int] + FILE_PATHS_FIELD_NUMBER: _ClassVar[int] id: str type: str scope: SegmentScope - topic: str collection: str metadata: UpdateMetadata - def __init__(self, id: _Optional[str] = ..., type: _Optional[str] = ..., scope: _Optional[_Union[SegmentScope, str]] = ..., topic: _Optional[str] = ..., collection: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ...) -> None: ... + file_paths: _containers.MessageMap[str, FilePaths] + def __init__(self, id: _Optional[str] = ..., type: _Optional[str] = ..., scope: _Optional[_Union[SegmentScope, str]] = ..., collection: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., file_paths: _Optional[_Mapping[str, FilePaths]] = ...) -> None: ... class Collection(_message.Message): - __slots__ = ["id", "name", "topic", "metadata", "dimension", "tenant", "database"] + __slots__ = ["id", "name", "metadata", "dimension", "tenant", "database", "log_position", "version"] ID_FIELD_NUMBER: _ClassVar[int] NAME_FIELD_NUMBER: _ClassVar[int] - TOPIC_FIELD_NUMBER: _ClassVar[int] METADATA_FIELD_NUMBER: _ClassVar[int] DIMENSION_FIELD_NUMBER: _ClassVar[int] TENANT_FIELD_NUMBER: _ClassVar[int] DATABASE_FIELD_NUMBER: _ClassVar[int] + LOG_POSITION_FIELD_NUMBER: _ClassVar[int] + VERSION_FIELD_NUMBER: _ClassVar[int] id: str name: str - topic: str metadata: UpdateMetadata dimension: int tenant: str database: str - def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., topic: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., dimension: _Optional[int] = ..., tenant: _Optional[str] = ..., database: _Optional[str] = ...) -> None: ... + log_position: int + version: int + def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., dimension: _Optional[int] = ..., tenant: _Optional[str] = ..., database: _Optional[str] = ..., log_position: _Optional[int] = ..., version: _Optional[int] = ...) -> None: ... class Database(_message.Message): __slots__ = ["id", "name", "tenant"] @@ -128,41 +137,35 @@ class UpdateMetadata(_message.Message): metadata: _containers.MessageMap[str, UpdateMetadataValue] def __init__(self, metadata: _Optional[_Mapping[str, UpdateMetadataValue]] = ...) -> None: ... -class SubmitEmbeddingRecord(_message.Message): - __slots__ = ["id", "vector", "metadata", "operation", "collection_id"] +class OperationRecord(_message.Message): + __slots__ = ["id", "vector", "metadata", "operation"] ID_FIELD_NUMBER: _ClassVar[int] VECTOR_FIELD_NUMBER: _ClassVar[int] METADATA_FIELD_NUMBER: _ClassVar[int] OPERATION_FIELD_NUMBER: _ClassVar[int] - COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] id: str vector: Vector metadata: UpdateMetadata operation: Operation - collection_id: str - def __init__(self, id: _Optional[str] = ..., vector: _Optional[_Union[Vector, _Mapping]] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., operation: _Optional[_Union[Operation, str]] = ..., collection_id: _Optional[str] = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., vector: _Optional[_Union[Vector, _Mapping]] = ..., metadata: _Optional[_Union[UpdateMetadata, _Mapping]] = ..., operation: _Optional[_Union[Operation, str]] = ...) -> None: ... class VectorEmbeddingRecord(_message.Message): - __slots__ = ["id", "seq_id", "vector"] + __slots__ = ["id", "vector"] ID_FIELD_NUMBER: _ClassVar[int] - SEQ_ID_FIELD_NUMBER: _ClassVar[int] VECTOR_FIELD_NUMBER: _ClassVar[int] id: str - seq_id: bytes vector: Vector - def __init__(self, id: _Optional[str] = ..., seq_id: _Optional[bytes] = ..., vector: _Optional[_Union[Vector, _Mapping]] = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., vector: _Optional[_Union[Vector, _Mapping]] = ...) -> None: ... class VectorQueryResult(_message.Message): - __slots__ = ["id", "seq_id", "distance", "vector"] + __slots__ = ["id", "distance", "vector"] ID_FIELD_NUMBER: _ClassVar[int] - SEQ_ID_FIELD_NUMBER: _ClassVar[int] DISTANCE_FIELD_NUMBER: _ClassVar[int] VECTOR_FIELD_NUMBER: _ClassVar[int] id: str - seq_id: bytes distance: float vector: Vector - def __init__(self, id: _Optional[str] = ..., seq_id: _Optional[bytes] = ..., distance: _Optional[float] = ..., vector: _Optional[_Union[Vector, _Mapping]] = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., distance: _Optional[float] = ..., vector: _Optional[_Union[Vector, _Mapping]] = ...) -> None: ... class VectorQueryResults(_message.Message): __slots__ = ["results"] diff --git a/chromadb/proto/chroma_pb2_grpc.py b/chromadb/proto/chroma_pb2_grpc.py index ccd53e449c0..1ce9a18b4cc 100644 --- a/chromadb/proto/chroma_pb2_grpc.py +++ b/chromadb/proto/chroma_pb2_grpc.py @@ -6,7 +6,9 @@ class VectorReaderStub(object): - """Vector Reader Interface""" + """Vector Reader Interface + + """ def __init__(self, channel): """Constructor. @@ -15,110 +17,89 @@ def __init__(self, channel): channel: A grpc.Channel. """ self.GetVectors = channel.unary_unary( - "/chroma.VectorReader/GetVectors", - request_serializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsResponse.FromString, - ) + '/chroma.VectorReader/GetVectors', + request_serializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsResponse.FromString, + ) self.QueryVectors = channel.unary_unary( - "/chroma.VectorReader/QueryVectors", - request_serializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsResponse.FromString, - ) + '/chroma.VectorReader/QueryVectors', + request_serializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsResponse.FromString, + ) class VectorReaderServicer(object): - """Vector Reader Interface""" + """Vector Reader Interface + + """ def GetVectors(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def QueryVectors(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def add_VectorReaderServicer_to_server(servicer, server): rpc_method_handlers = { - "GetVectors": grpc.unary_unary_rpc_method_handler( - servicer.GetVectors, - request_deserializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsResponse.SerializeToString, - ), - "QueryVectors": grpc.unary_unary_rpc_method_handler( - servicer.QueryVectors, - request_deserializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsResponse.SerializeToString, - ), + 'GetVectors': grpc.unary_unary_rpc_method_handler( + servicer.GetVectors, + request_deserializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsRequest.FromString, + response_serializer=chromadb_dot_proto_dot_chroma__pb2.GetVectorsResponse.SerializeToString, + ), + 'QueryVectors': grpc.unary_unary_rpc_method_handler( + servicer.QueryVectors, + request_deserializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsRequest.FromString, + response_serializer=chromadb_dot_proto_dot_chroma__pb2.QueryVectorsResponse.SerializeToString, + ), } generic_handler = grpc.method_handlers_generic_handler( - "chroma.VectorReader", rpc_method_handlers - ) + 'chroma.VectorReader', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) -# This class is part of an EXPERIMENTAL API. + # This class is part of an EXPERIMENTAL API. class VectorReader(object): - """Vector Reader Interface""" + """Vector Reader Interface + + """ @staticmethod - def GetVectors( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def GetVectors(request, target, - "/chroma.VectorReader/GetVectors", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.VectorReader/GetVectors', chromadb_dot_proto_dot_chroma__pb2.GetVectorsRequest.SerializeToString, chromadb_dot_proto_dot_chroma__pb2.GetVectorsResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def QueryVectors( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def QueryVectors(request, target, - "/chroma.VectorReader/QueryVectors", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.VectorReader/QueryVectors', chromadb_dot_proto_dot_chroma__pb2.QueryVectorsRequest.SerializeToString, chromadb_dot_proto_dot_chroma__pb2.QueryVectorsResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/chromadb/proto/convert.py b/chromadb/proto/convert.py index 78eaeb89101..c0fde0f58c8 100644 --- a/chromadb/proto/convert.py +++ b/chromadb/proto/convert.py @@ -3,17 +3,16 @@ from typing import Dict, Optional, Tuple, Union, cast from chromadb.api.types import Embedding import chromadb.proto.chroma_pb2 as proto -from chromadb.utils.messageid import bytes_to_int, int_to_bytes from chromadb.types import ( Collection, - EmbeddingRecord, + LogRecord, Metadata, Operation, ScalarEncoding, Segment, SegmentScope, SeqId, - SubmitEmbeddingRecord, + OperationRecord, UpdateMetadata, Vector, VectorEmbeddingRecord, @@ -112,17 +111,18 @@ def to_proto_update_metadata(metadata: UpdateMetadata) -> proto.UpdateMetadata: def from_proto_submit( - submit_embedding_record: proto.SubmitEmbeddingRecord, seq_id: SeqId -) -> EmbeddingRecord: - embedding, encoding = from_proto_vector(submit_embedding_record.vector) - record = EmbeddingRecord( - id=submit_embedding_record.id, - seq_id=seq_id, - embedding=embedding, - encoding=encoding, - metadata=from_proto_update_metadata(submit_embedding_record.metadata), - operation=from_proto_operation(submit_embedding_record.operation), - collection_id=UUID(hex=submit_embedding_record.collection_id), + operation_record: proto.OperationRecord, seq_id: SeqId +) -> LogRecord: + embedding, encoding = from_proto_vector(operation_record.vector) + record = LogRecord( + log_offset=seq_id, + record=OperationRecord( + id=operation_record.id, + embedding=embedding, + encoding=encoding, + metadata=from_proto_update_metadata(operation_record.metadata), + operation=from_proto_operation(operation_record.operation), + ), ) return record @@ -132,7 +132,6 @@ def from_proto_segment(segment: proto.Segment) -> Segment: id=UUID(hex=segment.id), type=segment.type, scope=from_proto_segment_scope(segment.scope), - topic=segment.topic if segment.HasField("topic") else None, collection=None if not segment.HasField("collection") else UUID(hex=segment.collection), @@ -147,7 +146,6 @@ def to_proto_segment(segment: Segment) -> proto.Segment: id=segment["id"].hex, type=segment["type"], scope=to_proto_segment_scope(segment["scope"]), - topic=segment["topic"], collection=None if segment["collection"] is None else segment["collection"].hex, metadata=None if segment["metadata"] is None @@ -195,7 +193,6 @@ def from_proto_collection(collection: proto.Collection) -> Collection: return Collection( id=UUID(hex=collection.id), name=collection.name, - topic=collection.topic, metadata=from_proto_metadata(collection.metadata) if collection.HasField("metadata") else None, @@ -211,7 +208,6 @@ def to_proto_collection(collection: Collection) -> proto.Collection: return proto.Collection( id=collection["id"].hex, name=collection["name"], - topic=collection["topic"], metadata=None if collection["metadata"] is None else to_proto_update_metadata(collection["metadata"]), @@ -238,8 +234,8 @@ def to_proto_operation(operation: Operation) -> proto.Operation: def to_proto_submit( - submit_record: SubmitEmbeddingRecord, -) -> proto.SubmitEmbeddingRecord: + submit_record: OperationRecord, +) -> proto.OperationRecord: vector = None if submit_record["embedding"] is not None and submit_record["encoding"] is not None: vector = to_proto_vector(submit_record["embedding"], submit_record["encoding"]) @@ -248,12 +244,11 @@ def to_proto_submit( if submit_record["metadata"] is not None: metadata = to_proto_update_metadata(submit_record["metadata"]) - return proto.SubmitEmbeddingRecord( + return proto.OperationRecord( id=submit_record["id"], vector=vector, metadata=metadata, operation=to_proto_operation(submit_record["operation"]), - collection_id=submit_record["collection_id"].hex, ) @@ -262,7 +257,6 @@ def from_proto_vector_embedding_record( ) -> VectorEmbeddingRecord: return VectorEmbeddingRecord( id=embedding_record.id, - seq_id=from_proto_seq_id(embedding_record.seq_id), embedding=from_proto_vector(embedding_record.vector)[0], ) @@ -273,7 +267,6 @@ def to_proto_vector_embedding_record( ) -> proto.VectorEmbeddingRecord: return proto.VectorEmbeddingRecord( id=embedding_record["id"], - seq_id=to_proto_seq_id(embedding_record["seq_id"]), vector=to_proto_vector(embedding_record["embedding"], encoding), ) @@ -283,15 +276,6 @@ def from_proto_vector_query_result( ) -> VectorQueryResult: return VectorQueryResult( id=vector_query_result.id, - seq_id=from_proto_seq_id(vector_query_result.seq_id), distance=vector_query_result.distance, embedding=from_proto_vector(vector_query_result.vector)[0], ) - - -def to_proto_seq_id(seq_id: SeqId) -> bytes: - return int_to_bytes(seq_id) - - -def from_proto_seq_id(seq_id: bytes) -> SeqId: - return bytes_to_int(seq_id) diff --git a/chromadb/proto/coordinator_pb2.py b/chromadb/proto/coordinator_pb2.py index fda6a099867..fde4981b6c2 100644 --- a/chromadb/proto/coordinator_pb2.py +++ b/chromadb/proto/coordinator_pb2.py @@ -15,48 +15,84 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n chromadb/proto/coordinator.proto\x12\x06\x63hroma\x1a\x1b\x63hromadb/proto/chroma.proto\x1a\x1bgoogle/protobuf/empty.proto\"A\n\x15\x43reateDatabaseRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"2\n\x12GetDatabaseRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\"Y\n\x13GetDatabaseResponse\x12\"\n\x08\x64\x61tabase\x18\x01 \x01(\x0b\x32\x10.chroma.Database\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"#\n\x13\x43reateTenantRequest\x12\x0c\n\x04name\x18\x02 \x01(\t\" \n\x10GetTenantRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\"S\n\x11GetTenantResponse\x12\x1e\n\x06tenant\x18\x01 \x01(\x0b\x32\x0e.chroma.Tenant\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"8\n\x14\x43reateSegmentRequest\x12 \n\x07segment\x18\x01 \x01(\x0b\x32\x0f.chroma.Segment\"\"\n\x14\x44\x65leteSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\"\xc2\x01\n\x12GetSegmentsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04type\x18\x02 \x01(\tH\x01\x88\x01\x01\x12(\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScopeH\x02\x88\x01\x01\x12\x12\n\x05topic\x18\x04 \x01(\tH\x03\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x04\x88\x01\x01\x42\x05\n\x03_idB\x07\n\x05_typeB\x08\n\x06_scopeB\x08\n\x06_topicB\r\n\x0b_collection\"X\n\x13GetSegmentsResponse\x12!\n\x08segments\x18\x01 \x03(\x0b\x32\x0f.chroma.Segment\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xfa\x01\n\x14UpdateSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0f\n\x05topic\x18\x02 \x01(\tH\x00\x12\x15\n\x0breset_topic\x18\x03 \x01(\x08H\x00\x12\x14\n\ncollection\x18\x04 \x01(\tH\x01\x12\x1a\n\x10reset_collection\x18\x05 \x01(\x08H\x01\x12*\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x02\x12\x18\n\x0ereset_metadata\x18\x07 \x01(\x08H\x02\x42\x0e\n\x0ctopic_updateB\x13\n\x11\x63ollection_updateB\x11\n\x0fmetadata_update\"\xe5\x01\n\x17\x43reateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rget_or_create\x18\x05 \x01(\x08H\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\tB\x0b\n\t_metadataB\x0c\n\n_dimensionB\x10\n\x0e_get_or_create\"s\n\x18\x43reateCollectionResponse\x12&\n\ncollection\x18\x01 \x01(\x0b\x32\x12.chroma.Collection\x12\x0f\n\x07\x63reated\x18\x02 \x01(\x08\x12\x1e\n\x06status\x18\x03 \x01(\x0b\x32\x0e.chroma.Status\"G\n\x17\x44\x65leteCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x03 \x01(\t\"\x8b\x01\n\x15GetCollectionsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05topic\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x04 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x05 \x01(\tB\x05\n\x03_idB\x07\n\x05_nameB\x08\n\x06_topic\"a\n\x16GetCollectionsResponse\x12\'\n\x0b\x63ollections\x18\x01 \x03(\x0b\x32\x12.chroma.Collection\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xde\x01\n\x17UpdateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x12\n\x05topic\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04name\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12*\n\x08metadata\x18\x05 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x12\x18\n\x0ereset_metadata\x18\x06 \x01(\x08H\x00\x42\x11\n\x0fmetadata_updateB\x08\n\x06_topicB\x07\n\x05_nameB\x0c\n\n_dimension2\xd6\x07\n\x05SysDB\x12I\n\x0e\x43reateDatabase\x12\x1d.chroma.CreateDatabaseRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12H\n\x0bGetDatabase\x12\x1a.chroma.GetDatabaseRequest\x1a\x1b.chroma.GetDatabaseResponse\"\x00\x12\x45\n\x0c\x43reateTenant\x12\x1b.chroma.CreateTenantRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12\x42\n\tGetTenant\x12\x18.chroma.GetTenantRequest\x1a\x19.chroma.GetTenantResponse\"\x00\x12G\n\rCreateSegment\x12\x1c.chroma.CreateSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12G\n\rDeleteSegment\x12\x1c.chroma.DeleteSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12H\n\x0bGetSegments\x12\x1a.chroma.GetSegmentsRequest\x1a\x1b.chroma.GetSegmentsResponse\"\x00\x12G\n\rUpdateSegment\x12\x1c.chroma.UpdateSegmentRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12W\n\x10\x43reateCollection\x12\x1f.chroma.CreateCollectionRequest\x1a .chroma.CreateCollectionResponse\"\x00\x12M\n\x10\x44\x65leteCollection\x12\x1f.chroma.DeleteCollectionRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12Q\n\x0eGetCollections\x12\x1d.chroma.GetCollectionsRequest\x1a\x1e.chroma.GetCollectionsResponse\"\x00\x12M\n\x10UpdateCollection\x12\x1f.chroma.UpdateCollectionRequest\x1a\x16.chroma.ChromaResponse\"\x00\x12>\n\nResetState\x12\x16.google.protobuf.Empty\x1a\x16.chroma.ChromaResponse\"\x00\x42\x43ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpbb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n chromadb/proto/coordinator.proto\x12\x06\x63hroma\x1a\x1b\x63hromadb/proto/chroma.proto\x1a\x1bgoogle/protobuf/empty.proto\"A\n\x15\x43reateDatabaseRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0e\n\x06tenant\x18\x03 \x01(\t\"8\n\x16\x43reateDatabaseResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"2\n\x12GetDatabaseRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\"Y\n\x13GetDatabaseResponse\x12\"\n\x08\x64\x61tabase\x18\x01 \x01(\x0b\x32\x10.chroma.Database\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"#\n\x13\x43reateTenantRequest\x12\x0c\n\x04name\x18\x02 \x01(\t\"6\n\x14\x43reateTenantResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\" \n\x10GetTenantRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\"S\n\x11GetTenantResponse\x12\x1e\n\x06tenant\x18\x01 \x01(\x0b\x32\x0e.chroma.Tenant\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"8\n\x14\x43reateSegmentRequest\x12 \n\x07segment\x18\x01 \x01(\x0b\x32\x0f.chroma.Segment\"7\n\x15\x43reateSegmentResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"\"\n\x14\x44\x65leteSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\"7\n\x15\x44\x65leteSegmentResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"\xa4\x01\n\x12GetSegmentsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04type\x18\x02 \x01(\tH\x01\x88\x01\x01\x12(\n\x05scope\x18\x03 \x01(\x0e\x32\x14.chroma.SegmentScopeH\x02\x88\x01\x01\x12\x17\n\ncollection\x18\x05 \x01(\tH\x03\x88\x01\x01\x42\x05\n\x03_idB\x07\n\x05_typeB\x08\n\x06_scopeB\r\n\x0b_collection\"X\n\x13GetSegmentsResponse\x12!\n\x08segments\x18\x01 \x03(\x0b\x32\x0f.chroma.Segment\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xc2\x01\n\x14UpdateSegmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x14\n\ncollection\x18\x04 \x01(\tH\x00\x12\x1a\n\x10reset_collection\x18\x05 \x01(\x08H\x00\x12*\n\x08metadata\x18\x06 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x01\x12\x18\n\x0ereset_metadata\x18\x07 \x01(\x08H\x01\x42\x13\n\x11\x63ollection_updateB\x11\n\x0fmetadata_update\"7\n\x15UpdateSegmentResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"\xe5\x01\n\x17\x43reateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12-\n\x08metadata\x18\x03 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rget_or_create\x18\x05 \x01(\x08H\x02\x88\x01\x01\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x07 \x01(\tB\x0b\n\t_metadataB\x0c\n\n_dimensionB\x10\n\x0e_get_or_create\"s\n\x18\x43reateCollectionResponse\x12&\n\ncollection\x18\x01 \x01(\x0b\x32\x12.chroma.Collection\x12\x0f\n\x07\x63reated\x18\x02 \x01(\x08\x12\x1e\n\x06status\x18\x03 \x01(\x0b\x32\x0e.chroma.Status\"G\n\x17\x44\x65leteCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0e\n\x06tenant\x18\x02 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x03 \x01(\t\":\n\x18\x44\x65leteCollectionResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"m\n\x15GetCollectionsRequest\x12\x0f\n\x02id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x11\n\x04name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x0e\n\x06tenant\x18\x04 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x05 \x01(\tB\x05\n\x03_idB\x07\n\x05_name\"a\n\x16GetCollectionsResponse\x12\'\n\x0b\x63ollections\x18\x01 \x03(\x0b\x32\x12.chroma.Collection\x12\x1e\n\x06status\x18\x02 \x01(\x0b\x32\x0e.chroma.Status\"\xc0\x01\n\x17UpdateCollectionRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\x04name\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tdimension\x18\x04 \x01(\x05H\x02\x88\x01\x01\x12*\n\x08metadata\x18\x05 \x01(\x0b\x32\x16.chroma.UpdateMetadataH\x00\x12\x18\n\x0ereset_metadata\x18\x06 \x01(\x08H\x00\x42\x11\n\x0fmetadata_updateB\x07\n\x05_nameB\x0c\n\n_dimension\":\n\x18UpdateCollectionResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\"O\n\x0cNotification\x12\n\n\x02id\x18\x01 \x01(\x03\x12\x15\n\rcollection_id\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x0e\n\x06status\x18\x04 \x01(\t\"4\n\x12ResetStateResponse\x12\x1e\n\x06status\x18\x01 \x01(\x0b\x32\x0e.chroma.Status\":\n%GetLastCompactionTimeForTenantRequest\x12\x11\n\ttenant_id\x18\x01 \x03(\t\"K\n\x18TenantLastCompactionTime\x12\x11\n\ttenant_id\x18\x01 \x01(\t\x12\x1c\n\x14last_compaction_time\x18\x02 \x01(\x03\"o\n&GetLastCompactionTimeForTenantResponse\x12\x45\n\x1btenant_last_compaction_time\x18\x01 \x03(\x0b\x32 .chroma.TenantLastCompactionTime\"n\n%SetLastCompactionTimeForTenantRequest\x12\x45\n\x1btenant_last_compaction_time\x18\x01 \x01(\x0b\x32 .chroma.TenantLastCompactionTime\"\xbc\x01\n\x1a\x46lushSegmentCompactionInfo\x12\x12\n\nsegment_id\x18\x01 \x01(\t\x12\x45\n\nfile_paths\x18\x02 \x03(\x0b\x32\x31.chroma.FlushSegmentCompactionInfo.FilePathsEntry\x1a\x43\n\x0e\x46ilePathsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12 \n\x05value\x18\x02 \x01(\x0b\x32\x11.chroma.FilePaths:\x02\x38\x01\"\xc3\x01\n FlushCollectionCompactionRequest\x12\x11\n\ttenant_id\x18\x01 \x01(\t\x12\x15\n\rcollection_id\x18\x02 \x01(\t\x12\x14\n\x0clog_position\x18\x03 \x01(\x03\x12\x1a\n\x12\x63ollection_version\x18\x04 \x01(\x05\x12\x43\n\x17segment_compaction_info\x18\x05 \x03(\x0b\x32\".chroma.FlushSegmentCompactionInfo\"t\n!FlushCollectionCompactionResponse\x12\x15\n\rcollection_id\x18\x01 \x01(\t\x12\x1a\n\x12\x63ollection_version\x18\x02 \x01(\x05\x12\x1c\n\x14last_compaction_time\x18\x03 \x01(\x03\x32\xf4\n\n\x05SysDB\x12Q\n\x0e\x43reateDatabase\x12\x1d.chroma.CreateDatabaseRequest\x1a\x1e.chroma.CreateDatabaseResponse\"\x00\x12H\n\x0bGetDatabase\x12\x1a.chroma.GetDatabaseRequest\x1a\x1b.chroma.GetDatabaseResponse\"\x00\x12K\n\x0c\x43reateTenant\x12\x1b.chroma.CreateTenantRequest\x1a\x1c.chroma.CreateTenantResponse\"\x00\x12\x42\n\tGetTenant\x12\x18.chroma.GetTenantRequest\x1a\x19.chroma.GetTenantResponse\"\x00\x12N\n\rCreateSegment\x12\x1c.chroma.CreateSegmentRequest\x1a\x1d.chroma.CreateSegmentResponse\"\x00\x12N\n\rDeleteSegment\x12\x1c.chroma.DeleteSegmentRequest\x1a\x1d.chroma.DeleteSegmentResponse\"\x00\x12H\n\x0bGetSegments\x12\x1a.chroma.GetSegmentsRequest\x1a\x1b.chroma.GetSegmentsResponse\"\x00\x12N\n\rUpdateSegment\x12\x1c.chroma.UpdateSegmentRequest\x1a\x1d.chroma.UpdateSegmentResponse\"\x00\x12W\n\x10\x43reateCollection\x12\x1f.chroma.CreateCollectionRequest\x1a .chroma.CreateCollectionResponse\"\x00\x12W\n\x10\x44\x65leteCollection\x12\x1f.chroma.DeleteCollectionRequest\x1a .chroma.DeleteCollectionResponse\"\x00\x12Q\n\x0eGetCollections\x12\x1d.chroma.GetCollectionsRequest\x1a\x1e.chroma.GetCollectionsResponse\"\x00\x12W\n\x10UpdateCollection\x12\x1f.chroma.UpdateCollectionRequest\x1a .chroma.UpdateCollectionResponse\"\x00\x12\x42\n\nResetState\x12\x16.google.protobuf.Empty\x1a\x1a.chroma.ResetStateResponse\"\x00\x12\x81\x01\n\x1eGetLastCompactionTimeForTenant\x12-.chroma.GetLastCompactionTimeForTenantRequest\x1a..chroma.GetLastCompactionTimeForTenantResponse\"\x00\x12i\n\x1eSetLastCompactionTimeForTenant\x12-.chroma.SetLastCompactionTimeForTenantRequest\x1a\x16.google.protobuf.Empty\"\x00\x12r\n\x19\x46lushCollectionCompaction\x12(.chroma.FlushCollectionCompactionRequest\x1a).chroma.FlushCollectionCompactionResponse\"\x00\x42:Z8github.com/chroma-core/chroma/go/pkg/proto/coordinatorpbb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'chromadb.proto.coordinator_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: DESCRIPTOR._options = None - DESCRIPTOR._serialized_options = b'ZAgithub.com/chroma/chroma-coordinator/internal/proto/coordinatorpb' + DESCRIPTOR._serialized_options = b'Z8github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb' + _FLUSHSEGMENTCOMPACTIONINFO_FILEPATHSENTRY._options = None + _FLUSHSEGMENTCOMPACTIONINFO_FILEPATHSENTRY._serialized_options = b'8\001' _globals['_CREATEDATABASEREQUEST']._serialized_start=102 _globals['_CREATEDATABASEREQUEST']._serialized_end=167 - _globals['_GETDATABASEREQUEST']._serialized_start=169 - _globals['_GETDATABASEREQUEST']._serialized_end=219 - _globals['_GETDATABASERESPONSE']._serialized_start=221 - _globals['_GETDATABASERESPONSE']._serialized_end=310 - _globals['_CREATETENANTREQUEST']._serialized_start=312 - _globals['_CREATETENANTREQUEST']._serialized_end=347 - _globals['_GETTENANTREQUEST']._serialized_start=349 - _globals['_GETTENANTREQUEST']._serialized_end=381 - _globals['_GETTENANTRESPONSE']._serialized_start=383 - _globals['_GETTENANTRESPONSE']._serialized_end=466 - _globals['_CREATESEGMENTREQUEST']._serialized_start=468 - _globals['_CREATESEGMENTREQUEST']._serialized_end=524 - _globals['_DELETESEGMENTREQUEST']._serialized_start=526 - _globals['_DELETESEGMENTREQUEST']._serialized_end=560 - _globals['_GETSEGMENTSREQUEST']._serialized_start=563 - _globals['_GETSEGMENTSREQUEST']._serialized_end=757 - _globals['_GETSEGMENTSRESPONSE']._serialized_start=759 - _globals['_GETSEGMENTSRESPONSE']._serialized_end=847 - _globals['_UPDATESEGMENTREQUEST']._serialized_start=850 - _globals['_UPDATESEGMENTREQUEST']._serialized_end=1100 - _globals['_CREATECOLLECTIONREQUEST']._serialized_start=1103 - _globals['_CREATECOLLECTIONREQUEST']._serialized_end=1332 - _globals['_CREATECOLLECTIONRESPONSE']._serialized_start=1334 - _globals['_CREATECOLLECTIONRESPONSE']._serialized_end=1449 - _globals['_DELETECOLLECTIONREQUEST']._serialized_start=1451 - _globals['_DELETECOLLECTIONREQUEST']._serialized_end=1522 - _globals['_GETCOLLECTIONSREQUEST']._serialized_start=1525 - _globals['_GETCOLLECTIONSREQUEST']._serialized_end=1664 - _globals['_GETCOLLECTIONSRESPONSE']._serialized_start=1666 - _globals['_GETCOLLECTIONSRESPONSE']._serialized_end=1763 - _globals['_UPDATECOLLECTIONREQUEST']._serialized_start=1766 - _globals['_UPDATECOLLECTIONREQUEST']._serialized_end=1988 - _globals['_SYSDB']._serialized_start=1991 - _globals['_SYSDB']._serialized_end=2973 + _globals['_CREATEDATABASERESPONSE']._serialized_start=169 + _globals['_CREATEDATABASERESPONSE']._serialized_end=225 + _globals['_GETDATABASEREQUEST']._serialized_start=227 + _globals['_GETDATABASEREQUEST']._serialized_end=277 + _globals['_GETDATABASERESPONSE']._serialized_start=279 + _globals['_GETDATABASERESPONSE']._serialized_end=368 + _globals['_CREATETENANTREQUEST']._serialized_start=370 + _globals['_CREATETENANTREQUEST']._serialized_end=405 + _globals['_CREATETENANTRESPONSE']._serialized_start=407 + _globals['_CREATETENANTRESPONSE']._serialized_end=461 + _globals['_GETTENANTREQUEST']._serialized_start=463 + _globals['_GETTENANTREQUEST']._serialized_end=495 + _globals['_GETTENANTRESPONSE']._serialized_start=497 + _globals['_GETTENANTRESPONSE']._serialized_end=580 + _globals['_CREATESEGMENTREQUEST']._serialized_start=582 + _globals['_CREATESEGMENTREQUEST']._serialized_end=638 + _globals['_CREATESEGMENTRESPONSE']._serialized_start=640 + _globals['_CREATESEGMENTRESPONSE']._serialized_end=695 + _globals['_DELETESEGMENTREQUEST']._serialized_start=697 + _globals['_DELETESEGMENTREQUEST']._serialized_end=731 + _globals['_DELETESEGMENTRESPONSE']._serialized_start=733 + _globals['_DELETESEGMENTRESPONSE']._serialized_end=788 + _globals['_GETSEGMENTSREQUEST']._serialized_start=791 + _globals['_GETSEGMENTSREQUEST']._serialized_end=955 + _globals['_GETSEGMENTSRESPONSE']._serialized_start=957 + _globals['_GETSEGMENTSRESPONSE']._serialized_end=1045 + _globals['_UPDATESEGMENTREQUEST']._serialized_start=1048 + _globals['_UPDATESEGMENTREQUEST']._serialized_end=1242 + _globals['_UPDATESEGMENTRESPONSE']._serialized_start=1244 + _globals['_UPDATESEGMENTRESPONSE']._serialized_end=1299 + _globals['_CREATECOLLECTIONREQUEST']._serialized_start=1302 + _globals['_CREATECOLLECTIONREQUEST']._serialized_end=1531 + _globals['_CREATECOLLECTIONRESPONSE']._serialized_start=1533 + _globals['_CREATECOLLECTIONRESPONSE']._serialized_end=1648 + _globals['_DELETECOLLECTIONREQUEST']._serialized_start=1650 + _globals['_DELETECOLLECTIONREQUEST']._serialized_end=1721 + _globals['_DELETECOLLECTIONRESPONSE']._serialized_start=1723 + _globals['_DELETECOLLECTIONRESPONSE']._serialized_end=1781 + _globals['_GETCOLLECTIONSREQUEST']._serialized_start=1783 + _globals['_GETCOLLECTIONSREQUEST']._serialized_end=1892 + _globals['_GETCOLLECTIONSRESPONSE']._serialized_start=1894 + _globals['_GETCOLLECTIONSRESPONSE']._serialized_end=1991 + _globals['_UPDATECOLLECTIONREQUEST']._serialized_start=1994 + _globals['_UPDATECOLLECTIONREQUEST']._serialized_end=2186 + _globals['_UPDATECOLLECTIONRESPONSE']._serialized_start=2188 + _globals['_UPDATECOLLECTIONRESPONSE']._serialized_end=2246 + _globals['_NOTIFICATION']._serialized_start=2248 + _globals['_NOTIFICATION']._serialized_end=2327 + _globals['_RESETSTATERESPONSE']._serialized_start=2329 + _globals['_RESETSTATERESPONSE']._serialized_end=2381 + _globals['_GETLASTCOMPACTIONTIMEFORTENANTREQUEST']._serialized_start=2383 + _globals['_GETLASTCOMPACTIONTIMEFORTENANTREQUEST']._serialized_end=2441 + _globals['_TENANTLASTCOMPACTIONTIME']._serialized_start=2443 + _globals['_TENANTLASTCOMPACTIONTIME']._serialized_end=2518 + _globals['_GETLASTCOMPACTIONTIMEFORTENANTRESPONSE']._serialized_start=2520 + _globals['_GETLASTCOMPACTIONTIMEFORTENANTRESPONSE']._serialized_end=2631 + _globals['_SETLASTCOMPACTIONTIMEFORTENANTREQUEST']._serialized_start=2633 + _globals['_SETLASTCOMPACTIONTIMEFORTENANTREQUEST']._serialized_end=2743 + _globals['_FLUSHSEGMENTCOMPACTIONINFO']._serialized_start=2746 + _globals['_FLUSHSEGMENTCOMPACTIONINFO']._serialized_end=2934 + _globals['_FLUSHSEGMENTCOMPACTIONINFO_FILEPATHSENTRY']._serialized_start=2867 + _globals['_FLUSHSEGMENTCOMPACTIONINFO_FILEPATHSENTRY']._serialized_end=2934 + _globals['_FLUSHCOLLECTIONCOMPACTIONREQUEST']._serialized_start=2937 + _globals['_FLUSHCOLLECTIONCOMPACTIONREQUEST']._serialized_end=3132 + _globals['_FLUSHCOLLECTIONCOMPACTIONRESPONSE']._serialized_start=3134 + _globals['_FLUSHCOLLECTIONCOMPACTIONRESPONSE']._serialized_end=3250 + _globals['_SYSDB']._serialized_start=3253 + _globals['_SYSDB']._serialized_end=4649 # @@protoc_insertion_point(module_scope) diff --git a/chromadb/proto/coordinator_pb2.pyi b/chromadb/proto/coordinator_pb2.pyi index 81545e4e283..b00a5be9b79 100644 --- a/chromadb/proto/coordinator_pb2.pyi +++ b/chromadb/proto/coordinator_pb2.pyi @@ -17,6 +17,12 @@ class CreateDatabaseRequest(_message.Message): tenant: str def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., tenant: _Optional[str] = ...) -> None: ... +class CreateDatabaseResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + class GetDatabaseRequest(_message.Message): __slots__ = ["name", "tenant"] NAME_FIELD_NUMBER: _ClassVar[int] @@ -39,6 +45,12 @@ class CreateTenantRequest(_message.Message): name: str def __init__(self, name: _Optional[str] = ...) -> None: ... +class CreateTenantResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + class GetTenantRequest(_message.Message): __slots__ = ["name"] NAME_FIELD_NUMBER: _ClassVar[int] @@ -59,25 +71,35 @@ class CreateSegmentRequest(_message.Message): segment: _chroma_pb2.Segment def __init__(self, segment: _Optional[_Union[_chroma_pb2.Segment, _Mapping]] = ...) -> None: ... +class CreateSegmentResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + class DeleteSegmentRequest(_message.Message): __slots__ = ["id"] ID_FIELD_NUMBER: _ClassVar[int] id: str def __init__(self, id: _Optional[str] = ...) -> None: ... +class DeleteSegmentResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + class GetSegmentsRequest(_message.Message): - __slots__ = ["id", "type", "scope", "topic", "collection"] + __slots__ = ["id", "type", "scope", "collection"] ID_FIELD_NUMBER: _ClassVar[int] TYPE_FIELD_NUMBER: _ClassVar[int] SCOPE_FIELD_NUMBER: _ClassVar[int] - TOPIC_FIELD_NUMBER: _ClassVar[int] COLLECTION_FIELD_NUMBER: _ClassVar[int] id: str type: str scope: _chroma_pb2.SegmentScope - topic: str collection: str - def __init__(self, id: _Optional[str] = ..., type: _Optional[str] = ..., scope: _Optional[_Union[_chroma_pb2.SegmentScope, str]] = ..., topic: _Optional[str] = ..., collection: _Optional[str] = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., type: _Optional[str] = ..., scope: _Optional[_Union[_chroma_pb2.SegmentScope, str]] = ..., collection: _Optional[str] = ...) -> None: ... class GetSegmentsResponse(_message.Message): __slots__ = ["segments", "status"] @@ -88,22 +110,24 @@ class GetSegmentsResponse(_message.Message): def __init__(self, segments: _Optional[_Iterable[_Union[_chroma_pb2.Segment, _Mapping]]] = ..., status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... class UpdateSegmentRequest(_message.Message): - __slots__ = ["id", "topic", "reset_topic", "collection", "reset_collection", "metadata", "reset_metadata"] + __slots__ = ["id", "collection", "reset_collection", "metadata", "reset_metadata"] ID_FIELD_NUMBER: _ClassVar[int] - TOPIC_FIELD_NUMBER: _ClassVar[int] - RESET_TOPIC_FIELD_NUMBER: _ClassVar[int] COLLECTION_FIELD_NUMBER: _ClassVar[int] RESET_COLLECTION_FIELD_NUMBER: _ClassVar[int] METADATA_FIELD_NUMBER: _ClassVar[int] RESET_METADATA_FIELD_NUMBER: _ClassVar[int] id: str - topic: str - reset_topic: bool collection: str reset_collection: bool metadata: _chroma_pb2.UpdateMetadata reset_metadata: bool - def __init__(self, id: _Optional[str] = ..., topic: _Optional[str] = ..., reset_topic: bool = ..., collection: _Optional[str] = ..., reset_collection: bool = ..., metadata: _Optional[_Union[_chroma_pb2.UpdateMetadata, _Mapping]] = ..., reset_metadata: bool = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., collection: _Optional[str] = ..., reset_collection: bool = ..., metadata: _Optional[_Union[_chroma_pb2.UpdateMetadata, _Mapping]] = ..., reset_metadata: bool = ...) -> None: ... + +class UpdateSegmentResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... class CreateCollectionRequest(_message.Message): __slots__ = ["id", "name", "metadata", "dimension", "get_or_create", "tenant", "database"] @@ -143,19 +167,23 @@ class DeleteCollectionRequest(_message.Message): database: str def __init__(self, id: _Optional[str] = ..., tenant: _Optional[str] = ..., database: _Optional[str] = ...) -> None: ... +class DeleteCollectionResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + class GetCollectionsRequest(_message.Message): - __slots__ = ["id", "name", "topic", "tenant", "database"] + __slots__ = ["id", "name", "tenant", "database"] ID_FIELD_NUMBER: _ClassVar[int] NAME_FIELD_NUMBER: _ClassVar[int] - TOPIC_FIELD_NUMBER: _ClassVar[int] TENANT_FIELD_NUMBER: _ClassVar[int] DATABASE_FIELD_NUMBER: _ClassVar[int] id: str name: str - topic: str tenant: str database: str - def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., topic: _Optional[str] = ..., tenant: _Optional[str] = ..., database: _Optional[str] = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., tenant: _Optional[str] = ..., database: _Optional[str] = ...) -> None: ... class GetCollectionsResponse(_message.Message): __slots__ = ["collections", "status"] @@ -166,17 +194,104 @@ class GetCollectionsResponse(_message.Message): def __init__(self, collections: _Optional[_Iterable[_Union[_chroma_pb2.Collection, _Mapping]]] = ..., status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... class UpdateCollectionRequest(_message.Message): - __slots__ = ["id", "topic", "name", "dimension", "metadata", "reset_metadata"] + __slots__ = ["id", "name", "dimension", "metadata", "reset_metadata"] ID_FIELD_NUMBER: _ClassVar[int] - TOPIC_FIELD_NUMBER: _ClassVar[int] NAME_FIELD_NUMBER: _ClassVar[int] DIMENSION_FIELD_NUMBER: _ClassVar[int] METADATA_FIELD_NUMBER: _ClassVar[int] RESET_METADATA_FIELD_NUMBER: _ClassVar[int] id: str - topic: str name: str dimension: int metadata: _chroma_pb2.UpdateMetadata reset_metadata: bool - def __init__(self, id: _Optional[str] = ..., topic: _Optional[str] = ..., name: _Optional[str] = ..., dimension: _Optional[int] = ..., metadata: _Optional[_Union[_chroma_pb2.UpdateMetadata, _Mapping]] = ..., reset_metadata: bool = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., name: _Optional[str] = ..., dimension: _Optional[int] = ..., metadata: _Optional[_Union[_chroma_pb2.UpdateMetadata, _Mapping]] = ..., reset_metadata: bool = ...) -> None: ... + +class UpdateCollectionResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + +class Notification(_message.Message): + __slots__ = ["id", "collection_id", "type", "status"] + ID_FIELD_NUMBER: _ClassVar[int] + COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] + TYPE_FIELD_NUMBER: _ClassVar[int] + STATUS_FIELD_NUMBER: _ClassVar[int] + id: int + collection_id: str + type: str + status: str + def __init__(self, id: _Optional[int] = ..., collection_id: _Optional[str] = ..., type: _Optional[str] = ..., status: _Optional[str] = ...) -> None: ... + +class ResetStateResponse(_message.Message): + __slots__ = ["status"] + STATUS_FIELD_NUMBER: _ClassVar[int] + status: _chroma_pb2.Status + def __init__(self, status: _Optional[_Union[_chroma_pb2.Status, _Mapping]] = ...) -> None: ... + +class GetLastCompactionTimeForTenantRequest(_message.Message): + __slots__ = ["tenant_id"] + TENANT_ID_FIELD_NUMBER: _ClassVar[int] + tenant_id: _containers.RepeatedScalarFieldContainer[str] + def __init__(self, tenant_id: _Optional[_Iterable[str]] = ...) -> None: ... + +class TenantLastCompactionTime(_message.Message): + __slots__ = ["tenant_id", "last_compaction_time"] + TENANT_ID_FIELD_NUMBER: _ClassVar[int] + LAST_COMPACTION_TIME_FIELD_NUMBER: _ClassVar[int] + tenant_id: str + last_compaction_time: int + def __init__(self, tenant_id: _Optional[str] = ..., last_compaction_time: _Optional[int] = ...) -> None: ... + +class GetLastCompactionTimeForTenantResponse(_message.Message): + __slots__ = ["tenant_last_compaction_time"] + TENANT_LAST_COMPACTION_TIME_FIELD_NUMBER: _ClassVar[int] + tenant_last_compaction_time: _containers.RepeatedCompositeFieldContainer[TenantLastCompactionTime] + def __init__(self, tenant_last_compaction_time: _Optional[_Iterable[_Union[TenantLastCompactionTime, _Mapping]]] = ...) -> None: ... + +class SetLastCompactionTimeForTenantRequest(_message.Message): + __slots__ = ["tenant_last_compaction_time"] + TENANT_LAST_COMPACTION_TIME_FIELD_NUMBER: _ClassVar[int] + tenant_last_compaction_time: TenantLastCompactionTime + def __init__(self, tenant_last_compaction_time: _Optional[_Union[TenantLastCompactionTime, _Mapping]] = ...) -> None: ... + +class FlushSegmentCompactionInfo(_message.Message): + __slots__ = ["segment_id", "file_paths"] + class FilePathsEntry(_message.Message): + __slots__ = ["key", "value"] + KEY_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] + key: str + value: _chroma_pb2.FilePaths + def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[_chroma_pb2.FilePaths, _Mapping]] = ...) -> None: ... + SEGMENT_ID_FIELD_NUMBER: _ClassVar[int] + FILE_PATHS_FIELD_NUMBER: _ClassVar[int] + segment_id: str + file_paths: _containers.MessageMap[str, _chroma_pb2.FilePaths] + def __init__(self, segment_id: _Optional[str] = ..., file_paths: _Optional[_Mapping[str, _chroma_pb2.FilePaths]] = ...) -> None: ... + +class FlushCollectionCompactionRequest(_message.Message): + __slots__ = ["tenant_id", "collection_id", "log_position", "collection_version", "segment_compaction_info"] + TENANT_ID_FIELD_NUMBER: _ClassVar[int] + COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] + LOG_POSITION_FIELD_NUMBER: _ClassVar[int] + COLLECTION_VERSION_FIELD_NUMBER: _ClassVar[int] + SEGMENT_COMPACTION_INFO_FIELD_NUMBER: _ClassVar[int] + tenant_id: str + collection_id: str + log_position: int + collection_version: int + segment_compaction_info: _containers.RepeatedCompositeFieldContainer[FlushSegmentCompactionInfo] + def __init__(self, tenant_id: _Optional[str] = ..., collection_id: _Optional[str] = ..., log_position: _Optional[int] = ..., collection_version: _Optional[int] = ..., segment_compaction_info: _Optional[_Iterable[_Union[FlushSegmentCompactionInfo, _Mapping]]] = ...) -> None: ... + +class FlushCollectionCompactionResponse(_message.Message): + __slots__ = ["collection_id", "collection_version", "last_compaction_time"] + COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] + COLLECTION_VERSION_FIELD_NUMBER: _ClassVar[int] + LAST_COMPACTION_TIME_FIELD_NUMBER: _ClassVar[int] + collection_id: str + collection_version: int + last_compaction_time: int + def __init__(self, collection_id: _Optional[str] = ..., collection_version: _Optional[int] = ..., last_compaction_time: _Optional[int] = ...) -> None: ... diff --git a/chromadb/proto/coordinator_pb2_grpc.py b/chromadb/proto/coordinator_pb2_grpc.py index 117c568c715..92ede663915 100644 --- a/chromadb/proto/coordinator_pb2_grpc.py +++ b/chromadb/proto/coordinator_pb2_grpc.py @@ -2,7 +2,6 @@ """Client and server classes corresponding to protobuf-defined services.""" import grpc -from chromadb.proto import chroma_pb2 as chromadb_dot_proto_dot_chroma__pb2 from chromadb.proto import coordinator_pb2 as chromadb_dot_proto_dot_coordinator__pb2 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 @@ -17,70 +16,85 @@ def __init__(self, channel): channel: A grpc.Channel. """ self.CreateDatabase = channel.unary_unary( - "/chroma.SysDB/CreateDatabase", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/CreateDatabase', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseResponse.FromString, + ) self.GetDatabase = channel.unary_unary( - "/chroma.SysDB/GetDatabase", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseResponse.FromString, - ) + '/chroma.SysDB/GetDatabase', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseResponse.FromString, + ) self.CreateTenant = channel.unary_unary( - "/chroma.SysDB/CreateTenant", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateTenantRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/CreateTenant', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateTenantRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateTenantResponse.FromString, + ) self.GetTenant = channel.unary_unary( - "/chroma.SysDB/GetTenant", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantResponse.FromString, - ) + '/chroma.SysDB/GetTenant', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantResponse.FromString, + ) self.CreateSegment = channel.unary_unary( - "/chroma.SysDB/CreateSegment", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/CreateSegment', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentResponse.FromString, + ) self.DeleteSegment = channel.unary_unary( - "/chroma.SysDB/DeleteSegment", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/DeleteSegment', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentResponse.FromString, + ) self.GetSegments = channel.unary_unary( - "/chroma.SysDB/GetSegments", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsResponse.FromString, - ) + '/chroma.SysDB/GetSegments', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsResponse.FromString, + ) self.UpdateSegment = channel.unary_unary( - "/chroma.SysDB/UpdateSegment", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/UpdateSegment', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentResponse.FromString, + ) self.CreateCollection = channel.unary_unary( - "/chroma.SysDB/CreateCollection", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionResponse.FromString, - ) + '/chroma.SysDB/CreateCollection', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionResponse.FromString, + ) self.DeleteCollection = channel.unary_unary( - "/chroma.SysDB/DeleteCollection", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/DeleteCollection', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionResponse.FromString, + ) self.GetCollections = channel.unary_unary( - "/chroma.SysDB/GetCollections", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsResponse.FromString, - ) + '/chroma.SysDB/GetCollections', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsResponse.FromString, + ) self.UpdateCollection = channel.unary_unary( - "/chroma.SysDB/UpdateCollection", - request_serializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionRequest.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/UpdateCollection', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionResponse.FromString, + ) self.ResetState = channel.unary_unary( - "/chroma.SysDB/ResetState", - request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString, - response_deserializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - ) + '/chroma.SysDB/ResetState', + request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.ResetStateResponse.FromString, + ) + self.GetLastCompactionTimeForTenant = channel.unary_unary( + '/chroma.SysDB/GetLastCompactionTimeForTenant', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetLastCompactionTimeForTenantRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetLastCompactionTimeForTenantResponse.FromString, + ) + self.SetLastCompactionTimeForTenant = channel.unary_unary( + '/chroma.SysDB/SetLastCompactionTimeForTenant', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.SetLastCompactionTimeForTenantRequest.SerializeToString, + response_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString, + ) + self.FlushCollectionCompaction = channel.unary_unary( + '/chroma.SysDB/FlushCollectionCompaction', + request_serializer=chromadb_dot_proto_dot_coordinator__pb2.FlushCollectionCompactionRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_coordinator__pb2.FlushCollectionCompactionResponse.FromString, + ) class SysDBServicer(object): @@ -89,533 +103,460 @@ class SysDBServicer(object): def CreateDatabase(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def GetDatabase(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def CreateTenant(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def GetTenant(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def CreateSegment(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def DeleteSegment(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def GetSegments(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def UpdateSegment(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def CreateCollection(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def DeleteCollection(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def GetCollections(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def UpdateCollection(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def ResetState(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Method not implemented!") - raise NotImplementedError("Method not implemented!") + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetLastCompactionTimeForTenant(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def SetLastCompactionTimeForTenant(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def FlushCollectionCompaction(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') def add_SysDBServicer_to_server(servicer, server): rpc_method_handlers = { - "CreateDatabase": grpc.unary_unary_rpc_method_handler( - servicer.CreateDatabase, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "GetDatabase": grpc.unary_unary_rpc_method_handler( - servicer.GetDatabase, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseRequest.FromString, - response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseResponse.SerializeToString, - ), - "CreateTenant": grpc.unary_unary_rpc_method_handler( - servicer.CreateTenant, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateTenantRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "GetTenant": grpc.unary_unary_rpc_method_handler( - servicer.GetTenant, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantRequest.FromString, - response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantResponse.SerializeToString, - ), - "CreateSegment": grpc.unary_unary_rpc_method_handler( - servicer.CreateSegment, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "DeleteSegment": grpc.unary_unary_rpc_method_handler( - servicer.DeleteSegment, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "GetSegments": grpc.unary_unary_rpc_method_handler( - servicer.GetSegments, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsRequest.FromString, - response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsResponse.SerializeToString, - ), - "UpdateSegment": grpc.unary_unary_rpc_method_handler( - servicer.UpdateSegment, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "CreateCollection": grpc.unary_unary_rpc_method_handler( - servicer.CreateCollection, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionRequest.FromString, - response_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionResponse.SerializeToString, - ), - "DeleteCollection": grpc.unary_unary_rpc_method_handler( - servicer.DeleteCollection, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "GetCollections": grpc.unary_unary_rpc_method_handler( - servicer.GetCollections, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsRequest.FromString, - response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsResponse.SerializeToString, - ), - "UpdateCollection": grpc.unary_unary_rpc_method_handler( - servicer.UpdateCollection, - request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionRequest.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), - "ResetState": grpc.unary_unary_rpc_method_handler( - servicer.ResetState, - request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString, - response_serializer=chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.SerializeToString, - ), + 'CreateDatabase': grpc.unary_unary_rpc_method_handler( + servicer.CreateDatabase, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseResponse.SerializeToString, + ), + 'GetDatabase': grpc.unary_unary_rpc_method_handler( + servicer.GetDatabase, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseResponse.SerializeToString, + ), + 'CreateTenant': grpc.unary_unary_rpc_method_handler( + servicer.CreateTenant, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateTenantRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateTenantResponse.SerializeToString, + ), + 'GetTenant': grpc.unary_unary_rpc_method_handler( + servicer.GetTenant, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetTenantResponse.SerializeToString, + ), + 'CreateSegment': grpc.unary_unary_rpc_method_handler( + servicer.CreateSegment, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentResponse.SerializeToString, + ), + 'DeleteSegment': grpc.unary_unary_rpc_method_handler( + servicer.DeleteSegment, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentResponse.SerializeToString, + ), + 'GetSegments': grpc.unary_unary_rpc_method_handler( + servicer.GetSegments, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsResponse.SerializeToString, + ), + 'UpdateSegment': grpc.unary_unary_rpc_method_handler( + servicer.UpdateSegment, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentResponse.SerializeToString, + ), + 'CreateCollection': grpc.unary_unary_rpc_method_handler( + servicer.CreateCollection, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionResponse.SerializeToString, + ), + 'DeleteCollection': grpc.unary_unary_rpc_method_handler( + servicer.DeleteCollection, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionResponse.SerializeToString, + ), + 'GetCollections': grpc.unary_unary_rpc_method_handler( + servicer.GetCollections, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsResponse.SerializeToString, + ), + 'UpdateCollection': grpc.unary_unary_rpc_method_handler( + servicer.UpdateCollection, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionResponse.SerializeToString, + ), + 'ResetState': grpc.unary_unary_rpc_method_handler( + servicer.ResetState, + request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.ResetStateResponse.SerializeToString, + ), + 'GetLastCompactionTimeForTenant': grpc.unary_unary_rpc_method_handler( + servicer.GetLastCompactionTimeForTenant, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.GetLastCompactionTimeForTenantRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.GetLastCompactionTimeForTenantResponse.SerializeToString, + ), + 'SetLastCompactionTimeForTenant': grpc.unary_unary_rpc_method_handler( + servicer.SetLastCompactionTimeForTenant, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.SetLastCompactionTimeForTenantRequest.FromString, + response_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString, + ), + 'FlushCollectionCompaction': grpc.unary_unary_rpc_method_handler( + servicer.FlushCollectionCompaction, + request_deserializer=chromadb_dot_proto_dot_coordinator__pb2.FlushCollectionCompactionRequest.FromString, + response_serializer=chromadb_dot_proto_dot_coordinator__pb2.FlushCollectionCompactionResponse.SerializeToString, + ), } generic_handler = grpc.method_handlers_generic_handler( - "chroma.SysDB", rpc_method_handlers - ) + 'chroma.SysDB', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) -# This class is part of an EXPERIMENTAL API. + # This class is part of an EXPERIMENTAL API. class SysDB(object): """Missing associated documentation comment in .proto file.""" @staticmethod - def CreateDatabase( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def CreateDatabase(request, target, - "/chroma.SysDB/CreateDatabase", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/CreateDatabase', chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.CreateDatabaseResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def GetDatabase( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def GetDatabase(request, target, - "/chroma.SysDB/GetDatabase", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/GetDatabase', chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseRequest.SerializeToString, chromadb_dot_proto_dot_coordinator__pb2.GetDatabaseResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def CreateTenant( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def CreateTenant(request, target, - "/chroma.SysDB/CreateTenant", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/CreateTenant', chromadb_dot_proto_dot_coordinator__pb2.CreateTenantRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.CreateTenantResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def GetTenant( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def GetTenant(request, target, - "/chroma.SysDB/GetTenant", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/GetTenant', chromadb_dot_proto_dot_coordinator__pb2.GetTenantRequest.SerializeToString, chromadb_dot_proto_dot_coordinator__pb2.GetTenantResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def CreateSegment( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def CreateSegment(request, target, - "/chroma.SysDB/CreateSegment", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/CreateSegment', chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.CreateSegmentResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def DeleteSegment( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def DeleteSegment(request, target, - "/chroma.SysDB/DeleteSegment", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/DeleteSegment', chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.DeleteSegmentResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def GetSegments( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def GetSegments(request, target, - "/chroma.SysDB/GetSegments", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/GetSegments', chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsRequest.SerializeToString, chromadb_dot_proto_dot_coordinator__pb2.GetSegmentsResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def UpdateSegment( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def UpdateSegment(request, target, - "/chroma.SysDB/UpdateSegment", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/UpdateSegment', chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.UpdateSegmentResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def CreateCollection( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def CreateCollection(request, target, - "/chroma.SysDB/CreateCollection", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/CreateCollection', chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionRequest.SerializeToString, chromadb_dot_proto_dot_coordinator__pb2.CreateCollectionResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def DeleteCollection( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def DeleteCollection(request, target, - "/chroma.SysDB/DeleteCollection", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/DeleteCollection', chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.DeleteCollectionResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def GetCollections( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def GetCollections(request, target, - "/chroma.SysDB/GetCollections", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/GetCollections', chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsRequest.SerializeToString, chromadb_dot_proto_dot_coordinator__pb2.GetCollectionsResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def UpdateCollection( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def UpdateCollection(request, target, - "/chroma.SysDB/UpdateCollection", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/UpdateCollection', chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionRequest.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.UpdateCollectionResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod - def ResetState( - request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None, - ): - return grpc.experimental.unary_unary( - request, + def ResetState(request, target, - "/chroma.SysDB/ResetState", + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/ResetState', google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString, - chromadb_dot_proto_dot_chroma__pb2.ChromaResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - ) + chromadb_dot_proto_dot_coordinator__pb2.ResetStateResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def GetLastCompactionTimeForTenant(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/GetLastCompactionTimeForTenant', + chromadb_dot_proto_dot_coordinator__pb2.GetLastCompactionTimeForTenantRequest.SerializeToString, + chromadb_dot_proto_dot_coordinator__pb2.GetLastCompactionTimeForTenantResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def SetLastCompactionTimeForTenant(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/SetLastCompactionTimeForTenant', + chromadb_dot_proto_dot_coordinator__pb2.SetLastCompactionTimeForTenantRequest.SerializeToString, + google_dot_protobuf_dot_empty__pb2.Empty.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def FlushCollectionCompaction(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.SysDB/FlushCollectionCompaction', + chromadb_dot_proto_dot_coordinator__pb2.FlushCollectionCompactionRequest.SerializeToString, + chromadb_dot_proto_dot_coordinator__pb2.FlushCollectionCompactionResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/chromadb/proto/logservice_pb2.py b/chromadb/proto/logservice_pb2.py new file mode 100644 index 00000000000..eb8616ca46d --- /dev/null +++ b/chromadb/proto/logservice_pb2.py @@ -0,0 +1,50 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: chromadb/proto/logservice.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from chromadb.proto import chroma_pb2 as chromadb_dot_proto_dot_chroma__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1f\x63hromadb/proto/logservice.proto\x12\x06\x63hroma\x1a\x1b\x63hromadb/proto/chroma.proto"R\n\x0fPushLogsRequest\x12\x15\n\rcollection_id\x18\x01 \x01(\t\x12(\n\x07records\x18\x02 \x03(\x0b\x32\x17.chroma.OperationRecord"(\n\x10PushLogsResponse\x12\x14\n\x0crecord_count\x18\x01 \x01(\x05"n\n\x0fPullLogsRequest\x12\x15\n\rcollection_id\x18\x01 \x01(\t\x12\x19\n\x11start_from_offset\x18\x02 \x01(\x03\x12\x12\n\nbatch_size\x18\x03 \x01(\x05\x12\x15\n\rend_timestamp\x18\x04 \x01(\x03"H\n\tLogRecord\x12\x12\n\nlog_offset\x18\x01 \x01(\x03\x12\'\n\x06record\x18\x02 \x01(\x0b\x32\x17.chroma.OperationRecord"6\n\x10PullLogsResponse\x12"\n\x07records\x18\x01 \x03(\x0b\x32\x11.chroma.LogRecord"W\n\x0e\x43ollectionInfo\x12\x15\n\rcollection_id\x18\x01 \x01(\t\x12\x18\n\x10\x66irst_log_offset\x18\x02 \x01(\x03\x12\x14\n\x0c\x66irst_log_ts\x18\x03 \x01(\x03"&\n$GetAllCollectionInfoToCompactRequest"\\\n%GetAllCollectionInfoToCompactResponse\x12\x33\n\x13\x61ll_collection_info\x18\x01 \x03(\x0b\x32\x16.chroma.CollectionInfo2\x8e\x02\n\nLogService\x12?\n\x08PushLogs\x12\x17.chroma.PushLogsRequest\x1a\x18.chroma.PushLogsResponse"\x00\x12?\n\x08PullLogs\x12\x17.chroma.PullLogsRequest\x1a\x18.chroma.PullLogsResponse"\x00\x12~\n\x1dGetAllCollectionInfoToCompact\x12,.chroma.GetAllCollectionInfoToCompactRequest\x1a-.chroma.GetAllCollectionInfoToCompactResponse"\x00\x42\x39Z7github.com/chroma-core/chroma/go/pkg/proto/logservicepbb\x06proto3' +) + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages( + DESCRIPTOR, "chromadb.proto.logservice_pb2", _globals +) +if _descriptor._USE_C_DESCRIPTORS == False: + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"Z7github.com/chroma-core/chroma/go/pkg/proto/logservicepb" + ) + _globals["_PUSHLOGSREQUEST"]._serialized_start = 72 + _globals["_PUSHLOGSREQUEST"]._serialized_end = 154 + _globals["_PUSHLOGSRESPONSE"]._serialized_start = 156 + _globals["_PUSHLOGSRESPONSE"]._serialized_end = 196 + _globals["_PULLLOGSREQUEST"]._serialized_start = 198 + _globals["_PULLLOGSREQUEST"]._serialized_end = 308 + _globals["_LOGRECORD"]._serialized_start = 310 + _globals["_LOGRECORD"]._serialized_end = 382 + _globals["_PULLLOGSRESPONSE"]._serialized_start = 384 + _globals["_PULLLOGSRESPONSE"]._serialized_end = 438 + _globals["_COLLECTIONINFO"]._serialized_start = 440 + _globals["_COLLECTIONINFO"]._serialized_end = 527 + _globals["_GETALLCOLLECTIONINFOTOCOMPACTREQUEST"]._serialized_start = 529 + _globals["_GETALLCOLLECTIONINFOTOCOMPACTREQUEST"]._serialized_end = 567 + _globals["_GETALLCOLLECTIONINFOTOCOMPACTRESPONSE"]._serialized_start = 569 + _globals["_GETALLCOLLECTIONINFOTOCOMPACTRESPONSE"]._serialized_end = 661 + _globals["_LOGSERVICE"]._serialized_start = 664 + _globals["_LOGSERVICE"]._serialized_end = 934 +# @@protoc_insertion_point(module_scope) diff --git a/chromadb/proto/logservice_pb2.pyi b/chromadb/proto/logservice_pb2.pyi new file mode 100644 index 00000000000..7dce97aafd0 --- /dev/null +++ b/chromadb/proto/logservice_pb2.pyi @@ -0,0 +1,101 @@ +from chromadb.proto import chroma_pb2 as _chroma_pb2 +from google.protobuf.internal import containers as _containers +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from typing import ( + ClassVar as _ClassVar, + Iterable as _Iterable, + Mapping as _Mapping, + Optional as _Optional, + Union as _Union, +) + +DESCRIPTOR: _descriptor.FileDescriptor + +class PushLogsRequest(_message.Message): + __slots__ = ["collection_id", "records"] + COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] + RECORDS_FIELD_NUMBER: _ClassVar[int] + collection_id: str + records: _containers.RepeatedCompositeFieldContainer[_chroma_pb2.OperationRecord] + def __init__( + self, + collection_id: _Optional[str] = ..., + records: _Optional[ + _Iterable[_Union[_chroma_pb2.OperationRecord, _Mapping]] + ] = ..., + ) -> None: ... + +class PushLogsResponse(_message.Message): + __slots__ = ["record_count"] + RECORD_COUNT_FIELD_NUMBER: _ClassVar[int] + record_count: int + def __init__(self, record_count: _Optional[int] = ...) -> None: ... + +class PullLogsRequest(_message.Message): + __slots__ = ["collection_id", "start_from_offset", "batch_size", "end_timestamp"] + COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] + START_FROM_OFFSET_FIELD_NUMBER: _ClassVar[int] + BATCH_SIZE_FIELD_NUMBER: _ClassVar[int] + END_TIMESTAMP_FIELD_NUMBER: _ClassVar[int] + collection_id: str + start_from_offset: int + batch_size: int + end_timestamp: int + def __init__( + self, + collection_id: _Optional[str] = ..., + start_from_offset: _Optional[int] = ..., + batch_size: _Optional[int] = ..., + end_timestamp: _Optional[int] = ..., + ) -> None: ... + +class LogRecord(_message.Message): + __slots__ = ["log_offset", "record"] + LOG_OFFSET_FIELD_NUMBER: _ClassVar[int] + RECORD_FIELD_NUMBER: _ClassVar[int] + log_offset: int + record: _chroma_pb2.OperationRecord + def __init__( + self, + log_offset: _Optional[int] = ..., + record: _Optional[_Union[_chroma_pb2.OperationRecord, _Mapping]] = ..., + ) -> None: ... + +class PullLogsResponse(_message.Message): + __slots__ = ["records"] + RECORDS_FIELD_NUMBER: _ClassVar[int] + records: _containers.RepeatedCompositeFieldContainer[LogRecord] + def __init__( + self, records: _Optional[_Iterable[_Union[LogRecord, _Mapping]]] = ... + ) -> None: ... + +class CollectionInfo(_message.Message): + __slots__ = ["collection_id", "first_log_offset", "first_log_ts"] + COLLECTION_ID_FIELD_NUMBER: _ClassVar[int] + FIRST_LOG_OFFSET_FIELD_NUMBER: _ClassVar[int] + FIRST_LOG_TS_FIELD_NUMBER: _ClassVar[int] + collection_id: str + first_log_offset: int + first_log_ts: int + def __init__( + self, + collection_id: _Optional[str] = ..., + first_log_offset: _Optional[int] = ..., + first_log_ts: _Optional[int] = ..., + ) -> None: ... + +class GetAllCollectionInfoToCompactRequest(_message.Message): + __slots__ = [] + def __init__(self) -> None: ... + +class GetAllCollectionInfoToCompactResponse(_message.Message): + __slots__ = ["all_collection_info"] + ALL_COLLECTION_INFO_FIELD_NUMBER: _ClassVar[int] + all_collection_info: _containers.RepeatedCompositeFieldContainer[CollectionInfo] + def __init__( + self, + all_collection_info: _Optional[ + _Iterable[_Union[CollectionInfo, _Mapping]] + ] = ..., + ) -> None: ... diff --git a/chromadb/proto/logservice_pb2_grpc.py b/chromadb/proto/logservice_pb2_grpc.py new file mode 100644 index 00000000000..ab20441aa9a --- /dev/null +++ b/chromadb/proto/logservice_pb2_grpc.py @@ -0,0 +1,132 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +from chromadb.proto import logservice_pb2 as chromadb_dot_proto_dot_logservice__pb2 + + +class LogServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.PushLogs = channel.unary_unary( + '/chroma.LogService/PushLogs', + request_serializer=chromadb_dot_proto_dot_logservice__pb2.PushLogsRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_logservice__pb2.PushLogsResponse.FromString, + ) + self.PullLogs = channel.unary_unary( + '/chroma.LogService/PullLogs', + request_serializer=chromadb_dot_proto_dot_logservice__pb2.PullLogsRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_logservice__pb2.PullLogsResponse.FromString, + ) + self.GetAllCollectionInfoToCompact = channel.unary_unary( + '/chroma.LogService/GetAllCollectionInfoToCompact', + request_serializer=chromadb_dot_proto_dot_logservice__pb2.GetAllCollectionInfoToCompactRequest.SerializeToString, + response_deserializer=chromadb_dot_proto_dot_logservice__pb2.GetAllCollectionInfoToCompactResponse.FromString, + ) + + +class LogServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def PushLogs(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def PullLogs(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetAllCollectionInfoToCompact(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_LogServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'PushLogs': grpc.unary_unary_rpc_method_handler( + servicer.PushLogs, + request_deserializer=chromadb_dot_proto_dot_logservice__pb2.PushLogsRequest.FromString, + response_serializer=chromadb_dot_proto_dot_logservice__pb2.PushLogsResponse.SerializeToString, + ), + 'PullLogs': grpc.unary_unary_rpc_method_handler( + servicer.PullLogs, + request_deserializer=chromadb_dot_proto_dot_logservice__pb2.PullLogsRequest.FromString, + response_serializer=chromadb_dot_proto_dot_logservice__pb2.PullLogsResponse.SerializeToString, + ), + 'GetAllCollectionInfoToCompact': grpc.unary_unary_rpc_method_handler( + servicer.GetAllCollectionInfoToCompact, + request_deserializer=chromadb_dot_proto_dot_logservice__pb2.GetAllCollectionInfoToCompactRequest.FromString, + response_serializer=chromadb_dot_proto_dot_logservice__pb2.GetAllCollectionInfoToCompactResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'chroma.LogService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class LogService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def PushLogs(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.LogService/PushLogs', + chromadb_dot_proto_dot_logservice__pb2.PushLogsRequest.SerializeToString, + chromadb_dot_proto_dot_logservice__pb2.PushLogsResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def PullLogs(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.LogService/PullLogs', + chromadb_dot_proto_dot_logservice__pb2.PullLogsRequest.SerializeToString, + chromadb_dot_proto_dot_logservice__pb2.PullLogsResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def GetAllCollectionInfoToCompact(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/chroma.LogService/GetAllCollectionInfoToCompact', + chromadb_dot_proto_dot_logservice__pb2.GetAllCollectionInfoToCompactRequest.SerializeToString, + chromadb_dot_proto_dot_logservice__pb2.GetAllCollectionInfoToCompactResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/chromadb/quota/__init__.py b/chromadb/quota/__init__.py new file mode 100644 index 00000000000..d74a369462d --- /dev/null +++ b/chromadb/quota/__init__.py @@ -0,0 +1,94 @@ +from abc import abstractmethod +from enum import Enum +from typing import Optional, Literal + +from chromadb import Documents, Embeddings +from chromadb.api import Metadatas +from chromadb.config import ( + Component, + System, +) + + +class Resource(Enum): + METADATA_KEY_LENGTH = "METADATA_KEY_LENGTH" + METADATA_VALUE_LENGTH = "METADATA_VALUE_LENGTH" + DOCUMENT_SIZE = "DOCUMENT_SIZE" + ADD_PER_MINUTE = "ADD_PER_MINUTE" + QUERY_PER_MINUTE = "QUERY_PER_MINUTE" + GET_PER_MINUTE = "QUERY_PER_MINUTE" + EMBEDDINGS_DIMENSION = "EMBEDDINGS_DIMENSION" + + +class QuotaError(Exception): + def __init__(self, resource: Resource, quota: int, actual: int): + super().__init__(f"quota error. resource: {resource} quota: {quota} actual: {actual}") + self.quota = quota + self.actual = actual + self.resource = resource + +class QuotaProvider(Component): + """ + Retrieves quotas for resources within a system. + + Methods: + get_for_subject(resource, subject=None, tier=None): + Returns the quota for a given resource, optionally considering the tier and subject. + """ + def __init__(self, system: System) -> None: + super().__init__(system) + self.system = system + + @abstractmethod + def get_for_subject(self, resource: Resource, subject: Optional[str] = None, tier: Optional[str] = None) -> \ + Optional[int]: + pass + + +class QuotaEnforcer(Component): + """ + Enforces quota restrictions on various resources using quota provider. + + Methods: + static_check(metadatas=None, documents=None, embeddings=None, collection_id=None): + Performs static checks against quotas for metadatas, documents, and embeddings. Raises QuotaError if limits are exceeded. + """ + def __init__(self, system: System) -> None: + super().__init__(system) + self.should_enforce = False + if system.settings.chroma_quota_provider_impl: + self._quota_provider = system.require(QuotaProvider) + self.should_enforce = True + self.system = system + + def static_check(self, metadatas: Optional[Metadatas] = None, documents: Optional[Documents] = None, + embeddings: Optional[Embeddings] = None, collection_id: Optional[str] = None): + + if not self.should_enforce: + return + metadata_key_length_quota = self._quota_provider.get_for_subject(resource=Resource.METADATA_KEY_LENGTH, + subject=collection_id) + metadata_value_length_quota = self._quota_provider.get_for_subject(resource=Resource.METADATA_VALUE_LENGTH, + subject=collection_id) + if metadatas and (metadata_key_length_quota or metadata_key_length_quota): + for metadata in metadatas: + for key in metadata: + if metadata_key_length_quota and len(key) > metadata_key_length_quota: + raise QuotaError(resource=Resource.METADATA_KEY_LENGTH, actual=len(key), + quota=metadata_key_length_quota) + if metadata_value_length_quota and isinstance(metadata[key], str) and len( + metadata[key]) > metadata_value_length_quota: + raise QuotaError(resource=Resource.METADATA_VALUE_LENGTH, actual=len(metadata[key]), + quota=metadata_value_length_quota) + document_size_quota = self._quota_provider.get_for_subject(resource=Resource.DOCUMENT_SIZE, subject=collection_id) + if document_size_quota and documents: + for document in documents: + if len(document) > document_size_quota: + raise QuotaError(resource=Resource.DOCUMENT_SIZE, actual=len(document), quota=document_size_quota) + embedding_dimension_quota = self._quota_provider.get_for_subject(resource=Resource.EMBEDDINGS_DIMENSION, + subject=collection_id) + if embedding_dimension_quota and embeddings: + for embedding in embeddings: + if len(embedding) > embedding_dimension_quota: + raise QuotaError(resource=Resource.EMBEDDINGS_DIMENSION, actual=len(embedding), + quota=embedding_dimension_quota) diff --git a/chromadb/quota/test_provider.py b/chromadb/quota/test_provider.py new file mode 100644 index 00000000000..484282fb7d0 --- /dev/null +++ b/chromadb/quota/test_provider.py @@ -0,0 +1,14 @@ +from typing import Optional + +from overrides import overrides + +from chromadb.quota import QuotaProvider, Resource + + +class QuotaProviderForTest(QuotaProvider): + def __init__(self, system) -> None: + super().__init__(system) + + @overrides + def get_for_subject(self, resource: Resource, subject: Optional[str] = "", tier: Optional[str] = "") -> Optional[int]: + pass diff --git a/chromadb/rate_limiting/__init__.py b/chromadb/rate_limiting/__init__.py new file mode 100644 index 00000000000..8bb2b7040e0 --- /dev/null +++ b/chromadb/rate_limiting/__init__.py @@ -0,0 +1,66 @@ +import inspect +from abc import abstractmethod +from functools import wraps +from typing import Optional, Any, Dict, Callable, cast + +from chromadb.config import Component +from chromadb.quota import QuotaProvider, Resource + + +class RateLimitError(Exception): + def __init__(self, resource: Resource, quota: int): + super().__init__(f"rate limit error. resource: {resource} quota: {quota}") + self.quota = quota + self.resource = resource + +class RateLimitingProvider(Component): + @abstractmethod + def is_allowed(self, key: str, quota: int, point: Optional[int] = 1) -> bool: + """ + Determines if a request identified by `key` can proceed given the current rate limit. + + :param key: The identifier for the requestor (unused in this simplified implementation). + :param quota: The quota which will be used for bucket size. + :param point: The number of tokens required to fulfill the request. + :return: True if the request can proceed, False otherwise. + """ + pass + + +def rate_limit( + subject: str, + resource: Resource +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + + def decorator(f: Callable[..., Any]) -> Callable[..., Any]: + args_name = inspect.getfullargspec(f)[0] + if subject not in args_name: + raise Exception(f'rate_limit decorator have unknown subject "{subject}", available {args_name}') + key_index = args_name.index(subject) + @wraps(f) + def wrapper(self, *args: Any, **kwargs: Dict[Any, Any]) -> Any: + # If not rate limiting provider is present, just run and return the function. + + if self._system.settings.chroma_rate_limiting_provider_impl is None: + return f(self, *args, **kwargs) + + if subject in kwargs: + subject_value = kwargs[subject] + else: + if len(args) < key_index: + return f(self, *args, **kwargs) + subject_value = args[key_index-1] + key_value = resource.value + "-" + str(subject_value) + self._system.settings.chroma_rate_limiting_provider_impl + quota_provider = self._system.require(QuotaProvider) + rate_limiter = self._system.require(RateLimitingProvider) + quota = quota_provider.get_for_subject(resource=resource,subject=str(subject_value)) + if quota is None: + return f(self, *args, **kwargs) + is_allowed = rate_limiter.is_allowed(key_value, quota) + if is_allowed is False: + raise RateLimitError(resource=resource.value, quota=quota) + return f(self, *args, **kwargs) + return wrapper + + return decorator \ No newline at end of file diff --git a/chromadb/rate_limiting/test_provider.py b/chromadb/rate_limiting/test_provider.py new file mode 100644 index 00000000000..6b97db3dad3 --- /dev/null +++ b/chromadb/rate_limiting/test_provider.py @@ -0,0 +1,15 @@ +from typing import Optional, Dict + +from overrides import overrides + +from chromadb.config import System +from chromadb.rate_limiting import RateLimitingProvider + + +class RateLimitingTestProvider(RateLimitingProvider): + def __init__(self, system: System): + super().__init__(system) + + @overrides + def is_allowed(self, key: str, quota: int, point: Optional[int] = 1) -> bool: + pass diff --git a/chromadb/segment/impl/__init__.py b/chromadb/segment/impl/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/chromadb/segment/impl/distributed/segment_directory.py b/chromadb/segment/impl/distributed/segment_directory.py index 70b766de675..d0738c054ee 100644 --- a/chromadb/segment/impl/distributed/segment_directory.py +++ b/chromadb/segment/impl/distributed/segment_directory.py @@ -10,6 +10,11 @@ from kubernetes import client, config, watch from kubernetes.client.rest import ApiException import threading +from chromadb.telemetry.opentelemetry import ( + OpenTelemetryGranularity, + add_attributes_to_current_span, + trace_method, +) from chromadb.utils.rendezvous_hash import assign, murmur3hasher @@ -226,6 +231,11 @@ def register_updated_segment_callback( ) -> None: raise NotImplementedError() + @trace_method( + "RendezvousHashSegmentDirectory._update_memberlist", + OpenTelemetryGranularity.ALL, + ) def _update_memberlist(self, memberlist: Memberlist) -> None: with self._curr_memberlist_mutex: + add_attributes_to_current_span({"new_memberlist": memberlist}) self._curr_memberlist = memberlist diff --git a/chromadb/segment/impl/distributed/server.py b/chromadb/segment/impl/distributed/server.py deleted file mode 100644 index d9a6c317f7a..00000000000 --- a/chromadb/segment/impl/distributed/server.py +++ /dev/null @@ -1,187 +0,0 @@ -from typing import Any, Dict, List, Sequence, Set -from uuid import UUID -from chromadb.config import Settings, System -from chromadb.ingest import CollectionAssignmentPolicy, Consumer -from chromadb.proto.chroma_pb2_grpc import ( - # SegmentServerServicer, - # add_SegmentServerServicer_to_server, - VectorReaderServicer, - add_VectorReaderServicer_to_server, -) -import chromadb.proto.chroma_pb2 as proto -import grpc -from concurrent import futures -from chromadb.proto.convert import ( - to_proto_vector_embedding_record -) -from chromadb.segment import SegmentImplementation, SegmentType -from chromadb.telemetry.opentelemetry import ( - OpenTelemetryClient -) -from chromadb.types import EmbeddingRecord -from chromadb.segment.distributed import MemberlistProvider, Memberlist -from chromadb.utils.rendezvous_hash import assign, murmur3hasher -from chromadb.ingest.impl.pulsar_admin import PulsarAdmin -import logging -import os - -# This file is a prototype. It will be replaced with a real distributed segment server -# written in a different language. This is just a proof of concept to get the distributed -# segment type working end to end. - -# Run this with python -m chromadb.segment.impl.distributed.server - -SEGMENT_TYPE_IMPLS = { - SegmentType.HNSW_DISTRIBUTED: "chromadb.segment.impl.vector.local_persistent_hnsw.PersistentLocalHnswSegment", -} - - -class SegmentServer(VectorReaderServicer): - _segment_cache: Dict[UUID, SegmentImplementation] = {} - _system: System - _opentelemetry_client: OpenTelemetryClient - _memberlist_provider: MemberlistProvider - _curr_memberlist: Memberlist - _assigned_topics: Set[str] - _topic_to_subscription: Dict[str, UUID] - _consumer: Consumer - - def __init__(self, system: System) -> None: - super().__init__() - self._system = system - - # Init dependency services - self._opentelemetry_client = system.require(OpenTelemetryClient) - # TODO: add term and epoch to segment server - self._memberlist_provider = system.require(MemberlistProvider) - self._memberlist_provider.set_memberlist_name("worker-memberlist") - self._assignment_policy = system.require(CollectionAssignmentPolicy) - self._create_pulsar_topics() - self._consumer = system.require(Consumer) - - # Init data - self._topic_to_subscription = {} - self._assigned_topics = set() - self._curr_memberlist = self._memberlist_provider.get_memberlist() - self._compute_assigned_topics() - - self._memberlist_provider.register_updated_memberlist_callback( - self._on_memberlist_update - ) - - def _compute_assigned_topics(self) -> None: - """Uses rendezvous hashing to compute the topics that this node is responsible for""" - if not self._curr_memberlist: - self._assigned_topics = set() - return - topics = self._assignment_policy.get_topics() - my_ip = os.environ["MY_POD_IP"] - new_assignments: List[str] = [] - for topic in topics: - assigned = assign(topic, self._curr_memberlist, murmur3hasher) - if assigned == my_ip: - new_assignments.append(topic) - new_assignments_set = set(new_assignments) - # TODO: We need to lock around this assignment - net_new_assignments = new_assignments_set - self._assigned_topics - removed_assignments = self._assigned_topics - new_assignments_set - - for topic in removed_assignments: - subscription = self._topic_to_subscription[topic] - self._consumer.unsubscribe(subscription) - del self._topic_to_subscription[topic] - - for topic in net_new_assignments: - subscription = self._consumer.subscribe(topic, self._on_message) - self._topic_to_subscription[topic] = subscription - - self._assigned_topics = new_assignments_set - print( - f"Topic assigment updated and now assigned to {len(self._assigned_topics)} topics" - ) - - def _on_memberlist_update(self, memberlist: Memberlist) -> None: - """Called when the memberlist is updated""" - self._curr_memberlist = memberlist - if len(self._curr_memberlist) > 0: - self._compute_assigned_topics() - else: - # In this case we'd want to warn that there are no members but - # this is not an error, as it could be that the cluster is just starting up - print("Memberlist is empty") - - def _on_message(self, embedding_records: Sequence[EmbeddingRecord]) -> None: - """Called when a message is received from the consumer""" - print(f"Received {len(embedding_records)} records") - print( - f"First record: {embedding_records[0]} is for collection {embedding_records[0]['collection_id']}" - ) - return None - - def _create_pulsar_topics(self) -> None: - """This creates the pulsar topics used by the system. - HACK: THIS IS COMPLETELY A HACK AND WILL BE REPLACED - BY A PROPER TOPIC MANAGEMENT SYSTEM IN THE COORDINATOR""" - topics = self._assignment_policy.get_topics() - admin = PulsarAdmin(self._system) - for topic in topics: - admin.create_topic(topic) - - def QueryVectors( - self, request: proto.QueryVectorsRequest, context: Any - ) -> proto.QueryVectorsResponse: - context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details("Query segment not implemented yet") - return proto.QueryVectorsResponse() - - # @trace_method( - # "SegmentServer.GetVectors", OpenTelemetryGranularity.OPERATION_AND_SEGMENT - # ) - # def GetVectors( - # self, request: proto.GetVectorsRequest, context: Any - # ) -> proto.GetVectorsResponse: - # segment_id = UUID(hex=request.segment_id) - # if segment_id not in self._segment_cache: - # context.set_code(grpc.StatusCode.NOT_FOUND) - # context.set_details("Segment not found") - # return proto.GetVectorsResponse() - # else: - # segment = self._segment_cache[segment_id] - # segment = cast(VectorReader, segment) - # segment_results = segment.get_vectors(request.ids) - # return_records = [] - # for record in segment_results: - # # TODO: encoding should be based on stored encoding for segment - # # For now we just assume float32 - # return_record = to_proto_vector_embedding_record( - # record, ScalarEncoding.FLOAT32 - # ) - # return_records.append(return_record) - # return proto.GetVectorsResponse(records=return_records) - - # def _cls(self, segment: Segment) -> Type[SegmentImplementation]: - # classname = SEGMENT_TYPE_IMPLS[SegmentType(segment["type"])] - # cls = get_class(classname, SegmentImplementation) - # return cls - - # def _create_instance(self, segment: Segment) -> None: - # if segment["id"] not in self._segment_cache: - # cls = self._cls(segment) - # instance = cls(self._system, segment) - # instance.start() - # self._segment_cache[segment["id"]] = instance - - -if __name__ == "__main__": - logging.basicConfig(level=logging.INFO) - system = System(Settings()) - server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) - segment_server = SegmentServer(system) - # add_SegmentServerServicer_to_server(segment_server, server) # type: ignore - add_VectorReaderServicer_to_server(segment_server, server) # type: ignore - server.add_insecure_port( - f"[::]:{system.settings.require('chroma_server_grpc_port')}" - ) - system.start() - server.start() - server.wait_for_termination() diff --git a/chromadb/segment/impl/manager/__init__.py b/chromadb/segment/impl/manager/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/chromadb/segment/impl/manager/cache/__init__.py b/chromadb/segment/impl/manager/cache/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/chromadb/segment/impl/manager/cache/cache.py b/chromadb/segment/impl/manager/cache/cache.py new file mode 100644 index 00000000000..80cab0d8e91 --- /dev/null +++ b/chromadb/segment/impl/manager/cache/cache.py @@ -0,0 +1,104 @@ +import uuid +from typing import Any, Callable +from chromadb.types import Segment +from overrides import override +from typing import Dict, Optional +from abc import ABC, abstractmethod + +class SegmentCache(ABC): + @abstractmethod + def get(self, key: uuid.UUID) -> Optional[Segment]: + pass + + @abstractmethod + def pop(self, key: uuid.UUID) -> Optional[Segment]: + pass + + @abstractmethod + def set(self, key: uuid.UUID, value: Segment) -> None: + pass + + @abstractmethod + def reset(self) -> None: + pass + + +class BasicCache(SegmentCache): + def __init__(self): + self.cache:Dict[uuid.UUID, Segment] = {} + + @override + def get(self, key: uuid.UUID) -> Optional[Segment]: + return self.cache.get(key) + + @override + def pop(self, key: uuid.UUID) -> Optional[Segment]: + return self.cache.pop(key, None) + + @override + def set(self, key: uuid.UUID, value: Segment) -> None: + self.cache[key] = value + + @override + def reset(self) -> None: + self.cache = {} + + +class SegmentLRUCache(BasicCache): + """A simple LRU cache implementation that handles objects with dynamic sizes. + The size of each object is determined by a user-provided size function.""" + + def __init__(self, capacity: int, size_func: Callable[[uuid.UUID], int], + callback: Optional[Callable[[uuid.UUID, Segment], Any]] = None): + self.capacity = capacity + self.size_func = size_func + self.cache: Dict[uuid.UUID, Segment] = {} + self.history = [] + self.callback = callback + + def _upsert_key(self, key: uuid.UUID): + if key in self.history: + self.history.remove(key) + self.history.append(key) + else: + self.history.append(key) + + @override + def get(self, key: uuid.UUID) -> Optional[Segment]: + self._upsert_key(key) + if key in self.cache: + return self.cache[key] + else: + return None + + @override + def pop(self, key: uuid.UUID) -> Optional[Segment]: + if key in self.history: + self.history.remove(key) + return self.cache.pop(key, None) + + + @override + def set(self, key: uuid.UUID, value: Segment) -> None: + if key in self.cache: + return + item_size = self.size_func(key) + key_sizes = {key: self.size_func(key) for key in self.cache} + total_size = sum(key_sizes.values()) + index = 0 + # Evict items if capacity is exceeded + while total_size + item_size > self.capacity and len(self.history) > index: + key_delete = self.history[index] + if key_delete in self.cache: + self.callback(key_delete, self.cache[key_delete]) + del self.cache[key_delete] + total_size -= key_sizes[key_delete] + index += 1 + + self.cache[key] = value + self._upsert_key(key) + + @override + def reset(self): + self.cache = {} + self.history = [] diff --git a/chromadb/segment/impl/manager/distributed.py b/chromadb/segment/impl/manager/distributed.py index c114b8a3c96..fb0352340c8 100644 --- a/chromadb/segment/impl/manager/distributed.py +++ b/chromadb/segment/impl/manager/distributed.py @@ -125,7 +125,7 @@ def hint_use_collection(self, collection_id: UUID, hint_type: Operation) -> None collection=collection_id, scope=SegmentScope.VECTOR ) known_types = set([k.value for k in SEGMENT_TYPE_IMPLS.keys()]) - segment = next(filter(lambda s: s["type"] in known_types, segments)) + segment = next(filter(lambda s: s["type"] in known_types, segments)) # noqa # grpc_url = self._segment_directory.get_segment_endpoint(segment) # if grpc_url not in self._segment_server_stubs: @@ -172,7 +172,6 @@ def _segment(type: SegmentType, scope: SegmentScope, collection: Collection) -> id=uuid4(), type=type.value, scope=scope, - topic=collection["topic"], collection=collection["id"], metadata=metadata, ) diff --git a/chromadb/segment/impl/manager/local.py b/chromadb/segment/impl/manager/local.py index 246d9a00c64..19710b921d8 100644 --- a/chromadb/segment/impl/manager/local.py +++ b/chromadb/segment/impl/manager/local.py @@ -7,6 +7,14 @@ VectorReader, S, ) +import logging +from chromadb.segment.impl.manager.cache.cache import ( + SegmentLRUCache, + BasicCache, + SegmentCache, +) +import os + from chromadb.config import System, get_class from chromadb.db.system import SysDB from overrides import override @@ -21,17 +29,17 @@ from chromadb.types import Collection, Operation, Segment, SegmentScope, Metadata from typing import Dict, Type, Sequence, Optional, cast from uuid import UUID, uuid4 -from collections import defaultdict import platform from chromadb.utils.lru_cache import LRUCache +from chromadb.utils.directory import get_directory_size + if platform.system() != "Windows": import resource elif platform.system() == "Windows": import ctypes - SEGMENT_TYPE_IMPLS = { SegmentType.SQLITE: "chromadb.segment.impl.metadata.sqlite.SqliteMetadataSegment", SegmentType.HNSW_LOCAL_MEMORY: "chromadb.segment.impl.vector.local_hnsw.LocalHnswSegment", @@ -47,9 +55,6 @@ class LocalSegmentManager(SegmentManager): _vector_instances_file_handle_cache: LRUCache[ UUID, PersistentLocalHnswSegment ] # LRU cache to manage file handles across vector segment instances - _segment_cache: Dict[ - UUID, Dict[SegmentScope, Segment] - ] # Tracks which segments are loaded for a given collection _vector_segment_type: SegmentType = SegmentType.HNSW_LOCAL_MEMORY _lock: Lock _max_file_handles: int @@ -59,8 +64,23 @@ def __init__(self, system: System): self._sysdb = self.require(SysDB) self._system = system self._opentelemetry_client = system.require(OpenTelemetryClient) + self.logger = logging.getLogger(__name__) self._instances = {} - self._segment_cache = defaultdict(dict) + self.segment_cache: Dict[SegmentScope, SegmentCache] = { + SegmentScope.METADATA: BasicCache() + } + if ( + system.settings.chroma_segment_cache_policy == "LRU" + and system.settings.chroma_memory_limit_bytes > 0 + ): + self.segment_cache[SegmentScope.VECTOR] = SegmentLRUCache( + capacity=system.settings.chroma_memory_limit_bytes, + callback=lambda k, v: self.callback_cache_evict(v), + size_func=lambda k: self._get_segment_disk_size(k), + ) + else: + self.segment_cache[SegmentScope.VECTOR] = BasicCache() + self._lock = Lock() # TODO: prototyping with distributed segment for now, but this should be a configurable option @@ -79,6 +99,13 @@ def __init__(self, system: System): segment_limit, callback=lambda _, v: v.close_persistent_index() ) + def callback_cache_evict(self, segment: Segment): + collection_id = segment["collection"] + self.logger.info(f"LRU cache evict collection {collection_id}") + instance = self._instance(segment) + instance.stop() + del self._instances[segment["id"]] + @override def start(self) -> None: for instance in self._instances.values(): @@ -97,7 +124,7 @@ def reset_state(self) -> None: instance.stop() instance.reset_state() self._instances = {} - self._segment_cache = defaultdict(dict) + self.segment_cache[SegmentScope.VECTOR].reset() super().reset_state() @trace_method( @@ -130,16 +157,38 @@ def delete_segments(self, collection_id: UUID) -> Sequence[UUID]: instance = self.get_segment(collection_id, MetadataReader) instance.delete() del self._instances[segment["id"]] - if collection_id in self._segment_cache: - if segment["scope"] in self._segment_cache[collection_id]: - del self._segment_cache[collection_id][segment["scope"]] - del self._segment_cache[collection_id] + if segment["scope"] is SegmentScope.VECTOR: + self.segment_cache[SegmentScope.VECTOR].pop(collection_id) + if segment["scope"] is SegmentScope.METADATA: + self.segment_cache[SegmentScope.METADATA].pop(collection_id) return [s["id"] for s in segments] @trace_method( "LocalSegmentManager.get_segment", OpenTelemetryGranularity.OPERATION_AND_SEGMENT, ) + def _get_segment_disk_size(self, collection_id: UUID) -> int: + segments = self._sysdb.get_segments( + collection=collection_id, scope=SegmentScope.VECTOR + ) + if len(segments) == 0: + return 0 + # With local segment manager (single server chroma), a collection always have one segment. + size = get_directory_size( + os.path.join( + self._system.settings.require("persist_directory"), + str(segments[0]["id"]), + ) + ) + return size + + def _get_segment_sysdb(self, collection_id: UUID, scope: SegmentScope): + segments = self._sysdb.get_segments(collection=collection_id, scope=scope) + known_types = set([k.value for k in SEGMENT_TYPE_IMPLS.keys()]) + # Get the first segment of a known type + segment = next(filter(lambda s: s["type"] in known_types, segments)) + return segment + @override def get_segment(self, collection_id: UUID, type: Type[S]) -> S: if type == MetadataReader: @@ -149,17 +198,15 @@ def get_segment(self, collection_id: UUID, type: Type[S]) -> S: else: raise ValueError(f"Invalid segment type: {type}") - if scope not in self._segment_cache[collection_id]: - segments = self._sysdb.get_segments(collection=collection_id, scope=scope) - known_types = set([k.value for k in SEGMENT_TYPE_IMPLS.keys()]) - # Get the first segment of a known type - segment = next(filter(lambda s: s["type"] in known_types, segments)) - self._segment_cache[collection_id][scope] = segment + segment = self.segment_cache[scope].get(collection_id) + if segment is None: + segment = self._get_segment_sysdb(collection_id, scope) + self.segment_cache[scope].set(collection_id, segment) # Instances must be atomically created, so we use a lock to ensure that only one thread # creates the instance. with self._lock: - instance = self._instance(self._segment_cache[collection_id][scope]) + instance = self._instance(segment) return cast(S, instance) @trace_method( @@ -206,7 +253,6 @@ def _segment(type: SegmentType, scope: SegmentScope, collection: Collection) -> id=uuid4(), type=type.value, scope=scope, - topic=collection["topic"], collection=collection["id"], metadata=metadata, ) diff --git a/chromadb/segment/impl/metadata/sqlite.py b/chromadb/segment/impl/metadata/sqlite.py index 7f6bae48813..23751c8954c 100644 --- a/chromadb/segment/impl/metadata/sqlite.py +++ b/chromadb/segment/impl/metadata/sqlite.py @@ -19,7 +19,7 @@ Where, WhereDocument, MetadataEmbeddingRecord, - EmbeddingRecord, + LogRecord, SeqId, Operation, UpdateMetadata, @@ -45,7 +45,7 @@ class SqliteMetadataSegment(MetadataReader): _db: SqliteDB _id: UUID _opentelemetry_client: OpenTelemetryClient - _topic: Optional[str] + _collection_id: Optional[UUID] _subscription: Optional[UUID] def __init__(self, system: System, segment: Segment): @@ -53,15 +53,17 @@ def __init__(self, system: System, segment: Segment): self._consumer = system.instance(Consumer) self._id = segment["id"] self._opentelemetry_client = system.require(OpenTelemetryClient) - self._topic = segment["topic"] + self._collection_id = segment["collection"] @trace_method("SqliteMetadataSegment.start", OpenTelemetryGranularity.ALL) @override def start(self) -> None: - if self._topic: + if self._collection_id: seq_id = self.max_seqid() self._subscription = self._consumer.subscribe( - self._topic, self._write_metadata, start=seq_id + collection_id=self._collection_id, + consume_fn=self._write_metadata, + start=seq_id, ) @trace_method("SqliteMetadataSegment.stop", OpenTelemetryGranularity.ALL) @@ -250,14 +252,11 @@ def _record(self, rows: Sequence[Tuple[Any, ...]]) -> MetadataEmbeddingRecord: return MetadataEmbeddingRecord( id=embedding_id, - seq_id=_decode_seq_id(seq_id), metadata=metadata or None, ) @trace_method("SqliteMetadataSegment._insert_record", OpenTelemetryGranularity.ALL) - def _insert_record( - self, cur: Cursor, record: EmbeddingRecord, upsert: bool - ) -> None: + def _insert_record(self, cur: Cursor, record: LogRecord, upsert: bool) -> None: """Add or update a single EmbeddingRecord into the DB""" t = Table("embeddings") @@ -266,11 +265,11 @@ def _insert_record( .into(t) .columns(t.segment_id, t.embedding_id, t.seq_id) .where(t.segment_id == ParameterValue(self._db.uuid_to_db(self._id))) - .where(t.embedding_id == ParameterValue(record["id"])) + .where(t.embedding_id == ParameterValue(record["operation_record"]["id"])) ).insert( ParameterValue(self._db.uuid_to_db(self._id)), - ParameterValue(record["id"]), - ParameterValue(_encode_seq_id(record["seq_id"])), + ParameterValue(record["operation_record"]["id"]), + ParameterValue(_encode_seq_id(record["log_offset"])), ) sql, params = get_sql(q) sql = sql + "RETURNING id" @@ -279,15 +278,18 @@ def _insert_record( except sqlite3.IntegrityError: # Can't use INSERT OR REPLACE here because it changes the primary key. if upsert: - return self._update_record(cur, record) + # Cast here because the OpenTel decorators obfuscate the type + return cast(None, self._update_record(cur, record)) else: - logger.warning(f"Insert of existing embedding ID: {record['id']}") + logger.warning( + f"Insert of existing embedding ID: {record['operation_record']['id']}" + ) # We are trying to add for a record that already exists. Fail the call. # We don't throw an exception since this is in principal an async path return - if record["metadata"]: - self._update_metadata(cur, id, record["metadata"]) + if record["operation_record"]["metadata"]: + self._update_metadata(cur, id, record["operation_record"]["metadata"]) @trace_method( "SqliteMetadataSegment._update_metadata", OpenTelemetryGranularity.ALL @@ -401,27 +403,51 @@ def insert_into_fulltext_search() -> None: insert_into_fulltext_search() @trace_method("SqliteMetadataSegment._delete_record", OpenTelemetryGranularity.ALL) - def _delete_record(self, cur: Cursor, record: EmbeddingRecord) -> None: + def _delete_record(self, cur: Cursor, record: LogRecord) -> None: """Delete a single EmbeddingRecord from the DB""" t = Table("embeddings") + fts_t = Table("embedding_fulltext_search") q = ( self._db.querybuilder() .from_(t) .where(t.segment_id == ParameterValue(self._db.uuid_to_db(self._id))) - .where(t.embedding_id == ParameterValue(record["id"])) + .where(t.embedding_id == ParameterValue(record["operation_record"]["id"])) + .delete() + ) + q_fts = ( + self._db.querybuilder() + .from_(fts_t) .delete() + .where( + fts_t.rowid.isin( + self._db.querybuilder() + .from_(t) + .select(t.id) + .where( + t.segment_id == ParameterValue(self._db.uuid_to_db(self._id)) + ) + .where( + t.embedding_id + == ParameterValue(record["operation_record"]["id"]) + ) + ) + ) ) + cur.execute(*get_sql(q_fts)) sql, params = get_sql(q) sql = sql + " RETURNING id" result = cur.execute(sql, params).fetchone() if result is None: - logger.warning(f"Delete of nonexisting embedding ID: {record['id']}") + logger.warning( + f"Delete of nonexisting embedding ID: {record['operation_record']['id']}" + ) else: id = result[0] # Manually delete metadata; cannot use cascade because # that triggers on replace metadata_t = Table("embedding_metadata") + q = ( self._db.querybuilder() .from_(metadata_t) @@ -432,28 +458,30 @@ def _delete_record(self, cur: Cursor, record: EmbeddingRecord) -> None: cur.execute(sql, params) @trace_method("SqliteMetadataSegment._update_record", OpenTelemetryGranularity.ALL) - def _update_record(self, cur: Cursor, record: EmbeddingRecord) -> None: + def _update_record(self, cur: Cursor, record: LogRecord) -> None: """Update a single EmbeddingRecord in the DB""" t = Table("embeddings") q = ( self._db.querybuilder() .update(t) - .set(t.seq_id, ParameterValue(_encode_seq_id(record["seq_id"]))) + .set(t.seq_id, ParameterValue(_encode_seq_id(record["log_offset"]))) .where(t.segment_id == ParameterValue(self._db.uuid_to_db(self._id))) - .where(t.embedding_id == ParameterValue(record["id"])) + .where(t.embedding_id == ParameterValue(record["operation_record"]["id"])) ) sql, params = get_sql(q) sql = sql + " RETURNING id" result = cur.execute(sql, params).fetchone() if result is None: - logger.warning(f"Update of nonexisting embedding ID: {record['id']}") + logger.warning( + f"Update of nonexisting embedding ID: {record['operation_record']['id']}" + ) else: id = result[0] - if record["metadata"]: - self._update_metadata(cur, id, record["metadata"]) + if record["operation_record"]["metadata"]: + self._update_metadata(cur, id, record["operation_record"]["metadata"]) @trace_method("SqliteMetadataSegment._write_metadata", OpenTelemetryGranularity.ALL) - def _write_metadata(self, records: Sequence[EmbeddingRecord]) -> None: + def _write_metadata(self, records: Sequence[LogRecord]) -> None: """Write embedding metadata to the database. Care should be taken to ensure records are append-only (that is, that seq-ids should increase monotonically)""" with self._db.tx() as cur: @@ -464,20 +492,20 @@ def _write_metadata(self, records: Sequence[EmbeddingRecord]) -> None: .columns("segment_id", "seq_id") .insert( ParameterValue(self._db.uuid_to_db(self._id)), - ParameterValue(_encode_seq_id(record["seq_id"])), + ParameterValue(_encode_seq_id(record["log_offset"])), ) ) sql, params = get_sql(q) sql = sql.replace("INSERT", "INSERT OR REPLACE") cur.execute(sql, params) - if record["operation"] == Operation.ADD: + if record["operation_record"]["operation"] == Operation.ADD: self._insert_record(cur, record, False) - elif record["operation"] == Operation.UPSERT: + elif record["operation_record"]["operation"] == Operation.UPSERT: self._insert_record(cur, record, True) - elif record["operation"] == Operation.DELETE: + elif record["operation_record"]["operation"] == Operation.DELETE: self._delete_record(cur, record) - elif record["operation"] == Operation.UPDATE: + elif record["operation_record"]["operation"] == Operation.UPDATE: self._update_record(cur, record) @trace_method( @@ -573,6 +601,7 @@ def _where_doc_criterion( def delete(self) -> None: t = Table("embeddings") t1 = Table("embedding_metadata") + t2 = Table("embedding_fulltext_search") q0 = ( self._db.querybuilder() .from_(t1) @@ -603,7 +632,23 @@ def delete(self) -> None: ) ) ) + q_fts = ( + self._db.querybuilder() + .from_(t2) + .delete() + .where( + t2.rowid.isin( + self._db.querybuilder() + .from_(t) + .select(t.id) + .where( + t.segment_id == ParameterValue(self._db.uuid_to_db(self._id)) + ) + ) + ) + ) with self._db.tx() as cur: + cur.execute(*get_sql(q_fts)) cur.execute(*get_sql(q0)) cur.execute(*get_sql(q)) diff --git a/chromadb/segment/impl/vector/batch.py b/chromadb/segment/impl/vector/batch.py index aac533b918f..43cbe886005 100644 --- a/chromadb/segment/impl/vector/batch.py +++ b/chromadb/segment/impl/vector/batch.py @@ -1,12 +1,12 @@ from typing import Dict, List, Set, cast -from chromadb.types import EmbeddingRecord, Operation, SeqId, Vector +from chromadb.types import LogRecord, Operation, SeqId, Vector class Batch: """Used to model the set of changes as an atomic operation""" - _ids_to_records: Dict[str, EmbeddingRecord] + _ids_to_records: Dict[str, LogRecord] _deleted_ids: Set[str] _written_ids: Set[str] _upsert_add_ids: Set[str] # IDs that are being added in an upsert @@ -37,9 +37,12 @@ def get_written_ids(self) -> List[str]: def get_written_vectors(self, ids: List[str]) -> List[Vector]: """Get the list of vectors to write in this batch""" - return [cast(Vector, self._ids_to_records[id]["embedding"]) for id in ids] + return [ + cast(Vector, self._ids_to_records[id]["operation_record"]["embedding"]) + for id in ids + ] - def get_record(self, id: str) -> EmbeddingRecord: + def get_record(self, id: str) -> LogRecord: """Get the record for a given ID""" return self._ids_to_records[id] @@ -51,24 +54,33 @@ def is_deleted(self, id: str) -> bool: def delete_count(self) -> int: return len(self._deleted_ids) - def apply(self, record: EmbeddingRecord, exists_already: bool = False) -> None: + def apply(self, record: LogRecord, exists_already: bool = False) -> None: """Apply an embedding record to this batch. Records passed to this method are assumed to be validated for correctness. For example, a delete or update presumes the ID exists in the index. An add presumes the ID does not exist in the index. The exists_already flag should be set to True if the ID does exist in the index, and False otherwise. """ - id = record["id"] - if record["operation"] == Operation.DELETE: + id = record["operation_record"]["id"] + if record["operation_record"]["operation"] == Operation.DELETE: # If the ID was previously written, remove it from the written set # And update the add/update/delete counts if id in self._written_ids: self._written_ids.remove(id) - if self._ids_to_records[id]["operation"] == Operation.ADD: + if ( + self._ids_to_records[id]["operation_record"]["operation"] + == Operation.ADD + ): self.add_count -= 1 - elif self._ids_to_records[id]["operation"] == Operation.UPDATE: + elif ( + self._ids_to_records[id]["operation_record"]["operation"] + == Operation.UPDATE + ): self.update_count -= 1 self._deleted_ids.add(id) - elif self._ids_to_records[id]["operation"] == Operation.UPSERT: + elif ( + self._ids_to_records[id]["operation_record"]["operation"] + == Operation.UPSERT + ): if id in self._upsert_add_ids: self.add_count -= 1 self._upsert_add_ids.remove(id) @@ -92,15 +104,15 @@ def apply(self, record: EmbeddingRecord, exists_already: bool = False) -> None: self._deleted_ids.remove(id) # Update the add/update counts - if record["operation"] == Operation.UPSERT: + if record["operation_record"]["operation"] == Operation.UPSERT: if not exists_already: self.add_count += 1 self._upsert_add_ids.add(id) else: self.update_count += 1 - elif record["operation"] == Operation.ADD: + elif record["operation_record"]["operation"] == Operation.ADD: self.add_count += 1 - elif record["operation"] == Operation.UPDATE: + elif record["operation_record"]["operation"] == Operation.UPDATE: self.update_count += 1 - self.max_seq_id = max(self.max_seq_id, record["seq_id"]) + self.max_seq_id = max(self.max_seq_id, record["log_offset"]) diff --git a/chromadb/segment/impl/vector/brute_force_index.py b/chromadb/segment/impl/vector/brute_force_index.py index f9466e3f3d4..3eef8e043d5 100644 --- a/chromadb/segment/impl/vector/brute_force_index.py +++ b/chromadb/segment/impl/vector/brute_force_index.py @@ -2,7 +2,7 @@ import numpy as np import numpy.typing as npt from chromadb.types import ( - EmbeddingRecord, + LogRecord, VectorEmbeddingRecord, VectorQuery, VectorQueryResult, @@ -59,7 +59,7 @@ def clear(self) -> None: self.free_indices = list(range(self.size)) self.vectors.fill(0) - def upsert(self, records: List[EmbeddingRecord]) -> None: + def upsert(self, records: List[LogRecord]) -> None: if len(records) + len(self) > self.size: raise Exception( "Index with capacity {} and {} current entries cannot add {} records".format( @@ -68,9 +68,9 @@ def upsert(self, records: List[EmbeddingRecord]) -> None: ) for i, record in enumerate(records): - id = record["id"] - vector = record["embedding"] - self.id_to_seq_id[id] = record["seq_id"] + id = record["operation_record"]["id"] + vector = record["operation_record"]["embedding"] + self.id_to_seq_id[id] = record["log_offset"] if id in self.deleted_ids: self.deleted_ids.remove(id) @@ -86,9 +86,9 @@ def upsert(self, records: List[EmbeddingRecord]) -> None: self.index_to_id[next_index] = id self.vectors[next_index] = vector - def delete(self, records: List[EmbeddingRecord]) -> None: + def delete(self, records: List[LogRecord]) -> None: for record in records: - id = record["id"] + id = record["operation_record"]["id"] if id in self.id_to_index: index = self.id_to_index[id] self.deleted_ids.add(id) @@ -113,7 +113,6 @@ def get_vectors( VectorEmbeddingRecord( id=id, embedding=self.vectors[self.id_to_index[id]].tolist(), - seq_id=self.id_to_seq_id[id], ) for id in target_ids ] @@ -145,7 +144,6 @@ def query(self, query: VectorQuery) -> Sequence[Sequence[VectorQueryResult]]: VectorQueryResult( id=id, distance=distances[i][j].item(), - seq_id=self.id_to_seq_id[id], embedding=self.vectors[j].tolist(), ) ) diff --git a/chromadb/segment/impl/vector/local_hnsw.py b/chromadb/segment/impl/vector/local_hnsw.py index e4437881b2a..b055762af4c 100644 --- a/chromadb/segment/impl/vector/local_hnsw.py +++ b/chromadb/segment/impl/vector/local_hnsw.py @@ -12,7 +12,7 @@ trace_method, ) from chromadb.types import ( - EmbeddingRecord, + LogRecord, VectorEmbeddingRecord, VectorQuery, VectorQueryResult, @@ -35,7 +35,7 @@ class LocalHnswSegment(VectorReader): _id: UUID _consumer: Consumer - _topic: Optional[str] + _collection: Optional[UUID] _subscription: UUID _settings: Settings _params: HnswParams @@ -49,6 +49,9 @@ class LocalHnswSegment(VectorReader): _id_to_label: Dict[str, int] _label_to_id: Dict[int, str] + # Note: As of the time of writing, this mapping is no longer needed. + # We merely keep it around for easy compatibility with the old code and + # debugging purposes. _id_to_seq_id: Dict[str, SeqId] _opentelemtry_client: OpenTelemetryClient @@ -56,7 +59,7 @@ class LocalHnswSegment(VectorReader): def __init__(self, system: System, segment: Segment): self._consumer = system.instance(Consumer) self._id = segment["id"] - self._topic = segment["topic"] + self._collection = segment["collection"] self._settings = system.settings self._params = HnswParams(segment["metadata"] or {}) @@ -84,10 +87,10 @@ def propagate_collection_metadata(metadata: Metadata) -> Optional[Metadata]: @override def start(self) -> None: super().start() - if self._topic: + if self._collection: seq_id = self.max_seqid() self._subscription = self._consumer.subscribe( - self._topic, self._write_records, start=seq_id + self._collection, self._write_records, start=seq_id ) @trace_method("LocalHnswSegment.stop", OpenTelemetryGranularity.ALL) @@ -116,10 +119,7 @@ def get_vectors( for label, vector in zip(labels, vectors): id = self._label_to_id[label] - seq_id = self._id_to_seq_id[id] - results.append( - VectorEmbeddingRecord(id=id, seq_id=seq_id, embedding=vector) - ) + results.append(VectorEmbeddingRecord(id=id, embedding=vector)) return results @@ -168,7 +168,6 @@ def filter_function(label: int) -> bool: result_labels[result_i], distances[result_i] ): id = self._label_to_id[label] - seq_id = self._id_to_seq_id[id] if query["include_embeddings"]: embedding = self._index.get_items([label])[0] else: @@ -176,7 +175,6 @@ def filter_function(label: int) -> bool: results.append( VectorQueryResult( id=id, - seq_id=seq_id, distance=distance.item(), embedding=embedding, ) @@ -272,7 +270,7 @@ def _apply_batch(self, batch: Batch) -> None: # If that succeeds, update the mappings for i, id in enumerate(written_ids): - self._id_to_seq_id[id] = batch.get_record(id)["seq_id"] + self._id_to_seq_id[id] = batch.get_record(id)["log_offset"] self._id_to_label[id] = labels_to_write[i] self._label_to_id[labels_to_write[i]] = id @@ -283,7 +281,7 @@ def _apply_batch(self, batch: Batch) -> None: self._max_seq_id = batch.max_seq_id @trace_method("LocalHnswSegment._write_records", OpenTelemetryGranularity.ALL) - def _write_records(self, records: Sequence[EmbeddingRecord]) -> None: + def _write_records(self, records: Sequence[LogRecord]) -> None: """Add a batch of embeddings to the index""" if not self._running: raise RuntimeError("Cannot add embeddings to stopped component") @@ -293,9 +291,9 @@ def _write_records(self, records: Sequence[EmbeddingRecord]) -> None: batch = Batch() for record in records: - self._max_seq_id = max(self._max_seq_id, record["seq_id"]) - id = record["id"] - op = record["operation"] + self._max_seq_id = max(self._max_seq_id, record["log_offset"]) + id = record["operation_record"]["id"] + op = record["operation_record"]["operation"] label = self._id_to_label.get(id, None) if op == Operation.DELETE: @@ -305,12 +303,12 @@ def _write_records(self, records: Sequence[EmbeddingRecord]) -> None: logger.warning(f"Delete of nonexisting embedding ID: {id}") elif op == Operation.UPDATE: - if record["embedding"] is not None: + if record["operation_record"]["embedding"] is not None: if label is not None: batch.apply(record) else: logger.warning( - f"Update of nonexisting embedding ID: {record['id']}" + f"Update of nonexisting embedding ID: {record['operation_record']['id']}" ) elif op == Operation.ADD: if not label: diff --git a/chromadb/segment/impl/vector/local_persistent_hnsw.py b/chromadb/segment/impl/vector/local_persistent_hnsw.py index 4ab60a1725d..4574ba3c954 100644 --- a/chromadb/segment/impl/vector/local_persistent_hnsw.py +++ b/chromadb/segment/impl/vector/local_persistent_hnsw.py @@ -17,7 +17,7 @@ trace_method, ) from chromadb.types import ( - EmbeddingRecord, + LogRecord, Metadata, Operation, Segment, @@ -222,14 +222,16 @@ def _apply_batch(self, batch: Batch) -> None: "PersistentLocalHnswSegment._write_records", OpenTelemetryGranularity.ALL ) @override - def _write_records(self, records: Sequence[EmbeddingRecord]) -> None: + def _write_records(self, records: Sequence[LogRecord]) -> None: """Add a batch of embeddings to the index""" if not self._running: raise RuntimeError("Cannot add embeddings to stopped component") with WriteRWLock(self._lock): for record in records: - if record["embedding"] is not None: - self._ensure_index(len(records), len(record["embedding"])) + if record["operation_record"]["embedding"] is not None: + self._ensure_index( + len(records), len(record["operation_record"]["embedding"]) + ) if not self._index_initialized: # If the index is not initialized here, it means that we have # not yet added any records to the index. So we can just @@ -237,9 +239,9 @@ def _write_records(self, records: Sequence[EmbeddingRecord]) -> None: continue self._brute_force_index = cast(BruteForceIndex, self._brute_force_index) - self._max_seq_id = max(self._max_seq_id, record["seq_id"]) - id = record["id"] - op = record["operation"] + self._max_seq_id = max(self._max_seq_id, record["log_offset"]) + id = record["operation_record"]["id"] + op = record["operation_record"]["operation"] exists_in_index = self._id_to_label.get( id, None ) is not None or self._brute_force_index.has_id(id) @@ -254,23 +256,23 @@ def _write_records(self, records: Sequence[EmbeddingRecord]) -> None: logger.warning(f"Delete of nonexisting embedding ID: {id}") elif op == Operation.UPDATE: - if record["embedding"] is not None: + if record["operation_record"]["embedding"] is not None: if exists_in_index: self._curr_batch.apply(record) self._brute_force_index.upsert([record]) else: logger.warning( - f"Update of nonexisting embedding ID: {record['id']}" + f"Update of nonexisting embedding ID: {record['operation_record']['id']}" ) elif op == Operation.ADD: - if record["embedding"] is not None: + if record["operation_record"]["embedding"] is not None: if not exists_in_index: self._curr_batch.apply(record, not exists_in_index) self._brute_force_index.upsert([record]) else: logger.warning(f"Add of existing embedding ID: {id}") elif op == Operation.UPSERT: - if record["embedding"] is not None: + if record["operation_record"]["embedding"] is not None: self._curr_batch.apply(record, exists_in_index) self._brute_force_index.upsert([record]) if len(self._curr_batch) >= self._batch_size: @@ -325,9 +327,8 @@ def get_vectors( for label, vector in zip(hnsw_labels, vectors): id = self._label_to_id[label] - seq_id = self._id_to_seq_id[id] results[id_to_index[id]] = VectorEmbeddingRecord( - id=id, seq_id=seq_id, embedding=vector + id=id, embedding=vector ) return results # type: ignore ## Python can't cast List with Optional to List with VectorEmbeddingRecord diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py index 529606a6c36..292f5038dea 100644 --- a/chromadb/server/fastapi/__init__.py +++ b/chromadb/server/fastapi/__init__.py @@ -35,6 +35,8 @@ InvalidDimensionException, InvalidHTTPVersion, ) +from chromadb.quota import QuotaError +from chromadb.rate_limiting import RateLimitError from chromadb.server.fastapi.types import ( AddEmbedding, CreateDatabase, @@ -50,7 +52,7 @@ import logging -from chromadb.server.fastapi.utils import fastapi_json_response, string_to_uuid as _uuid +from chromadb.utils.fastapi import fastapi_json_response, string_to_uuid as _uuid from chromadb.telemetry.opentelemetry.fastapi import instrument_fastapi from chromadb.types import Database, Tenant from chromadb.telemetry.product import ServerContext, ProductTelemetryClient @@ -140,6 +142,8 @@ def __init__(self, settings: Settings): allow_origins=settings.chroma_server_cors_allow_origins, allow_methods=["*"], ) + self._app.add_exception_handler(QuotaError, self.quota_exception_handler) + self._app.add_exception_handler(RateLimitError, self.rate_limit_exception_handler) self._app.on_event("shutdown")(self.shutdown) @@ -288,9 +292,22 @@ def shutdown(self) -> None: def app(self) -> fastapi.FastAPI: return self._app + async def rate_limit_exception_handler(self, request: Request, exc: RateLimitError): + return JSONResponse( + status_code=429, + content={"message": f"rate limit. resource: {exc.resource} quota: {exc.quota}"}, + ) + + def root(self) -> Dict[str, int]: return {"nanosecond heartbeat": self._api.heartbeat()} + async def quota_exception_handler(self, request: Request, exc: QuotaError): + return JSONResponse( + status_code=429, + content={"message": f"quota error. resource: {exc.resource} quota: {exc.quota} actual: {exc.actual}"}, + ) + def heartbeat(self) -> Dict[str, int]: return self.root() diff --git a/chromadb/telemetry/opentelemetry/grpc.py b/chromadb/telemetry/opentelemetry/grpc.py new file mode 100644 index 00000000000..175a342c56f --- /dev/null +++ b/chromadb/telemetry/opentelemetry/grpc.py @@ -0,0 +1,77 @@ +import binascii +import collections + +import grpc +from opentelemetry.trace import StatusCode, SpanKind + +from chromadb.telemetry.opentelemetry import tracer + + +class _ClientCallDetails( + collections.namedtuple('_ClientCallDetails', ('method', 'timeout', 'metadata', 'credentials')), + grpc.ClientCallDetails +): + pass + + +def _encode_span_id(span_id: int) -> str: + return binascii.hexlify(span_id.to_bytes(8, 'big')).decode() + + +def _encode_trace_id(trace_id: int) -> str: + return binascii.hexlify(trace_id.to_bytes(16, 'big')).decode() + +# Using OtelInterceptor with gRPC: +# 1. Instantiate the interceptor: interceptors = [OtelInterceptor()] +# 2. Intercept the channel: channel = grpc.intercept_channel(channel, *interceptors) + +class OtelInterceptor( + grpc.UnaryUnaryClientInterceptor, + grpc.UnaryStreamClientInterceptor, + grpc.StreamUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor +): + def _intercept_call(self, continuation, client_call_details, request_or_iterator): + if tracer is None: + return continuation(client_call_details, request_or_iterator) + with tracer.start_as_current_span(f"RPC {client_call_details.method}", kind=SpanKind.CLIENT) as span: + # Prepare metadata for propagation + metadata = client_call_details.metadata[:] if client_call_details.metadata else [] + metadata.extend([ + ("chroma-traceid", _encode_trace_id(span.get_span_context().trace_id)), + ("chroma-spanid", _encode_span_id(span.get_span_context().span_id)), + ]) + # Update client call details with new metadata + new_client_details = _ClientCallDetails( + client_call_details.method, + client_call_details.timeout, + tuple(metadata), # Ensure metadata is a tuple + client_call_details.credentials, + ) + try: + result = continuation(new_client_details, request_or_iterator) + # Set attributes based on the result + if hasattr(result, 'details') and result.details(): + span.set_attribute("rpc.detail", result.details()) + span.set_attribute("rpc.status_code", result.code().name.lower()) + # Set span status based on gRPC call result + if result.code() != grpc.StatusCode.OK: + span.set_status(StatusCode.ERROR, description=str(result.code())) + return result + except Exception as e: + # Log exception details and re-raise + span.set_attribute("rpc.error", str(e)) + span.set_status(StatusCode.ERROR, description=str(e)) + raise + + def intercept_unary_unary(self, continuation, client_call_details, request): + return self._intercept_call(continuation, client_call_details, request) + + def intercept_unary_stream(self, continuation, client_call_details, request): + return self._intercept_call(continuation, client_call_details, request) + + def intercept_stream_unary(self, continuation, client_call_details, request_iterator): + return self._intercept_call(continuation, client_call_details, request_iterator) + + def intercept_stream_stream(self, continuation, client_call_details, request_iterator): + return self._intercept_call(continuation, client_call_details, request_iterator) diff --git a/chromadb/test/client/create_http_client_with_basic_auth.py b/chromadb/test/client/create_http_client_with_basic_auth.py new file mode 100644 index 00000000000..2f884c571c1 --- /dev/null +++ b/chromadb/test/client/create_http_client_with_basic_auth.py @@ -0,0 +1,27 @@ +# This file is used by test_create_http_client.py to test the initialization +# of an HttpClient class with auth settings. +# +# See https://github.com/chroma-core/chroma/issues/1554 + +import chromadb +from chromadb.config import Settings +import sys + +def main() -> None: + try: + chromadb.HttpClient( + host='localhost', + port=8000, + settings=Settings( + chroma_client_auth_provider="chromadb.auth.basic.BasicAuthClientProvider", + chroma_client_auth_credentials="admin:testDb@home2" + ) + ) + except ValueError: + # We don't expect to be able to connect to Chroma. We just want to make sure + # there isn't an ImportError. + sys.exit(0) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/chromadb/test/client/test_cloud_client.py b/chromadb/test/client/test_cloud_client.py index aee869ca1c5..48b0252789b 100644 --- a/chromadb/test/client/test_cloud_client.py +++ b/chromadb/test/client/test_cloud_client.py @@ -61,7 +61,7 @@ def mock_cloud_server(valid_token: str) -> Generator[System, None, None]: settings = Settings( chroma_api_impl="chromadb.api.fastapi.FastAPI", chroma_server_host=TEST_CLOUD_HOST, - chroma_server_http_port=str(port), + chroma_server_http_port=port, chroma_client_auth_provider="chromadb.auth.token.TokenAuthClientProvider", chroma_client_auth_credentials=valid_token, chroma_client_auth_token_transport_header=TOKEN_TRANSPORT_HEADER, @@ -82,7 +82,7 @@ def test_valid_key(mock_cloud_server: System, valid_token: str) -> None: database=DEFAULT_DATABASE, api_key=valid_token, cloud_host=TEST_CLOUD_HOST, - cloud_port=mock_cloud_server.settings.chroma_server_http_port, # type: ignore + cloud_port=mock_cloud_server.settings.chroma_server_http_port or 8000, enable_ssl=False, ) @@ -98,7 +98,7 @@ def test_invalid_key(mock_cloud_server: System, valid_token: str) -> None: database=DEFAULT_DATABASE, api_key=invalid_token, cloud_host=TEST_CLOUD_HOST, - cloud_port=mock_cloud_server.settings.chroma_server_http_port, # type: ignore + cloud_port=mock_cloud_server.settings.chroma_server_http_port or 8000, enable_ssl=False, ) client.heartbeat() diff --git a/chromadb/test/client/test_create_http_client.py b/chromadb/test/client/test_create_http_client.py new file mode 100644 index 00000000000..293e5031ba9 --- /dev/null +++ b/chromadb/test/client/test_create_http_client.py @@ -0,0 +1,15 @@ +import subprocess + +# Needs to be a module, not a file, so that local imports work. +TEST_MODULE = "chromadb.test.client.create_http_client_with_basic_auth" + + +def test_main() -> None: + # This is the only way to test what we want to test: pytest does a bunch of + # importing and other module stuff in the background, so we need a clean + # python process to make sure we're not circular-importing. + # + # See https://github.com/chroma-core/chroma/issues/1554 + + res = subprocess.run(['python', '-m', TEST_MODULE]) + assert res.returncode == 0 diff --git a/chromadb/test/conftest.py b/chromadb/test/conftest.py index 087cb2271bd..f196c662747 100644 --- a/chromadb/test/conftest.py +++ b/chromadb/test/conftest.py @@ -3,6 +3,7 @@ import os import shutil import socket +import subprocess import tempfile import time from typing import ( @@ -16,6 +17,7 @@ Tuple, Callable, ) +from uuid import UUID import hypothesis import pytest @@ -28,7 +30,7 @@ from chromadb.config import Settings, System from chromadb.db.mixins import embeddings_queue from chromadb.ingest import Producer -from chromadb.types import SeqId, SubmitEmbeddingRecord +from chromadb.types import SeqId, OperationRecord from chromadb.api.client import Client as ClientCreator root_logger = logging.getLogger() @@ -47,7 +49,6 @@ ) hypothesis.settings.load_profile(os.getenv("HYPOTHESIS_PROFILE", "dev")) - NOT_CLUSTER_ONLY = os.getenv("CHROMA_CLUSTER_TEST_ONLY") != "1" @@ -58,6 +59,35 @@ def skip_if_not_cluster() -> pytest.MarkDecorator: ) +def generate_self_signed_certificate() -> None: + config_path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "openssl.cnf" + ) + print(f"Config path: {config_path}") # Debug print to verify path + if not os.path.exists(config_path): + raise FileNotFoundError(f"Config file not found at {config_path}") + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:4096", + "-keyout", + "serverkey.pem", + "-out", + "servercert.pem", + "-days", + "365", + "-nodes", + "-subj", + "/CN=localhost", + "-config", + config_path, + ] + ) + + def find_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) @@ -77,6 +107,8 @@ def _run_server( chroma_server_authz_provider: Optional[str] = None, chroma_server_authz_config_file: Optional[str] = None, chroma_server_authz_config: Optional[Dict[str, Any]] = None, + chroma_server_ssl_certfile: Optional[str] = None, + chroma_server_ssl_keyfile: Optional[str] = None, ) -> None: """Run a Chroma server locally""" if is_persistent and persist_directory: @@ -123,6 +155,8 @@ def _run_server( port=port, log_level="error", timeout_keep_alive=30, + ssl_keyfile=chroma_server_ssl_keyfile, + ssl_certfile=chroma_server_ssl_certfile, ) @@ -152,6 +186,8 @@ def _fastapi_fixture( chroma_server_authz_provider: Optional[str] = None, chroma_server_authz_config_file: Optional[str] = None, chroma_server_authz_config: Optional[Dict[str, Any]] = None, + chroma_server_ssl_certfile: Optional[str] = None, + chroma_server_ssl_keyfile: Optional[str] = None, ) -> Generator[System, None, None]: """Fixture generator that launches a server in a separate process, and yields a fastapi client connect to it""" @@ -171,6 +207,8 @@ def _fastapi_fixture( Optional[str], Optional[str], Optional[Dict[str, Any]], + Optional[str], + Optional[str], ] = ( port, False, @@ -183,6 +221,8 @@ def _fastapi_fixture( chroma_server_authz_provider, chroma_server_authz_config_file, chroma_server_authz_config, + chroma_server_ssl_certfile, + chroma_server_ssl_keyfile, ) persist_directory = None if is_persistent: @@ -199,17 +239,21 @@ def _fastapi_fixture( chroma_server_authz_provider, chroma_server_authz_config_file, chroma_server_authz_config, + chroma_server_ssl_certfile, + chroma_server_ssl_keyfile, ) proc = ctx.Process(target=_run_server, args=args, daemon=True) proc.start() settings = Settings( chroma_api_impl="chromadb.api.fastapi.FastAPI", chroma_server_host="localhost", - chroma_server_http_port=str(port), + chroma_server_http_port=port, allow_reset=True, chroma_client_auth_provider=chroma_client_auth_provider, chroma_client_auth_credentials=chroma_client_auth_credentials, chroma_client_auth_token_transport_header=chroma_client_auth_token_transport_header, + chroma_server_ssl_verify=chroma_server_ssl_certfile, + chroma_server_ssl_enabled=True if chroma_server_ssl_certfile else False, ) system = System(settings) api = system.instance(ServerAPI) @@ -231,10 +275,20 @@ def fastapi_persistent() -> Generator[System, None, None]: return _fastapi_fixture(is_persistent=True) +def fastapi_ssl() -> Generator[System, None, None]: + generate_self_signed_certificate() + return _fastapi_fixture( + is_persistent=False, + chroma_server_ssl_certfile="./servercert.pem", + chroma_server_ssl_keyfile="./serverkey.pem", + ) + + def basic_http_client() -> Generator[System, None, None]: settings = Settings( chroma_api_impl="chromadb.api.fastapi.FastAPI", - chroma_server_http_port="8000", + chroma_server_http_port=8000, + chroma_server_host="localhost", allow_reset=True, ) system = System(settings) @@ -400,6 +454,11 @@ def system_fixtures_wrong_auth() -> List[Callable[[], Generator[System, None, No return fixtures +def system_fixtures_ssl() -> List[Callable[[], Generator[System, None, None]]]: + fixtures = [fastapi_ssl] + return fixtures + + @pytest.fixture(scope="module", params=system_fixtures_wrong_auth()) def system_wrong_auth( request: pytest.FixtureRequest, @@ -412,6 +471,11 @@ def system(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield next(request.param()) +@pytest.fixture(scope="module", params=system_fixtures_ssl()) +def system_ssl(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: + yield next(request.param()) + + @pytest.fixture(scope="module", params=system_fixtures_auth()) def system_auth(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield next(request.param()) @@ -432,6 +496,14 @@ def client(system: System) -> Generator[ClientAPI, None, None]: client.clear_system_cache() +@pytest.fixture(scope="function") +def client_ssl(system_ssl: System) -> Generator[ClientAPI, None, None]: + system_ssl.reset_state() + client = ClientCreator.from_system(system_ssl) + yield client + client.clear_system_cache() + + @pytest.fixture(scope="function") def api_wrong_cred( system_wrong_auth: System, @@ -456,24 +528,24 @@ class ProducerFn(Protocol): def __call__( self, producer: Producer, - topic: str, - embeddings: Iterator[SubmitEmbeddingRecord], + collection_id: UUID, + embeddings: Iterator[OperationRecord], n: int, - ) -> Tuple[Sequence[SubmitEmbeddingRecord], Sequence[SeqId]]: + ) -> Tuple[Sequence[OperationRecord], Sequence[SeqId]]: ... def produce_n_single( producer: Producer, - topic: str, - embeddings: Iterator[SubmitEmbeddingRecord], + collection_id: UUID, + embeddings: Iterator[OperationRecord], n: int, -) -> Tuple[Sequence[SubmitEmbeddingRecord], Sequence[SeqId]]: +) -> Tuple[Sequence[OperationRecord], Sequence[SeqId]]: submitted_embeddings = [] seq_ids = [] for _ in range(n): e = next(embeddings) - seq_id = producer.submit_embedding(topic, e) + seq_id = producer.submit_embedding(collection_id, e) submitted_embeddings.append(e) seq_ids.append(seq_id) return submitted_embeddings, seq_ids @@ -481,16 +553,16 @@ def produce_n_single( def produce_n_batch( producer: Producer, - topic: str, - embeddings: Iterator[SubmitEmbeddingRecord], + collection_id: UUID, + embeddings: Iterator[OperationRecord], n: int, -) -> Tuple[Sequence[SubmitEmbeddingRecord], Sequence[SeqId]]: +) -> Tuple[Sequence[OperationRecord], Sequence[SeqId]]: submitted_embeddings = [] seq_ids: Sequence[SeqId] = [] for _ in range(n): e = next(embeddings) submitted_embeddings.append(e) - seq_ids = producer.submit_embeddings(topic, submitted_embeddings) + seq_ids = producer.submit_embeddings(collection_id, submitted_embeddings) return submitted_embeddings, seq_ids diff --git a/chromadb/test/db/test_system.py b/chromadb/test/db/test_system.py index 9971d81af93..673fc9a46c1 100644 --- a/chromadb/test/db/test_system.py +++ b/chromadb/test/db/test_system.py @@ -11,7 +11,6 @@ from chromadb.config import ( DEFAULT_DATABASE, DEFAULT_TENANT, - Component, System, Settings, ) @@ -20,22 +19,15 @@ from pytest import FixtureRequest import uuid -PULSAR_TENANT = "default" -PULSAR_NAMESPACE = "default" +TENANT = "default" +NAMESPACE = "default" # These are the sample collections that are used in the tests below. Tests can override # the fields as needed. - -# HACK: In order to get the real grpc tests passing, we need the topic to use rendezvous -# hashing. This is because the grpc tests use the real grpc sysdb server and the -# rendezvous hashing is done in the segment server. We don't have a easy way to parameterize -# the assignment policy in the grpc tests, so we just use rendezvous hashing for all tests. -# by harcoding the topic to what we expect rendezvous hashing to return with 16 topics. sample_collections = [ Collection( id=uuid.UUID(int=1), name="test_collection_1", - topic=f"persistent://{PULSAR_TENANT}/{PULSAR_NAMESPACE}/chroma_log_1", metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3}, dimension=128, database=DEFAULT_DATABASE, @@ -44,7 +36,6 @@ Collection( id=uuid.UUID(int=2), name="test_collection_2", - topic=f"persistent://{PULSAR_TENANT}/{PULSAR_NAMESPACE}/chroma_log_14", metadata={"test_str": "str2", "test_int": 2, "test_float": 2.3}, dimension=None, database=DEFAULT_DATABASE, @@ -53,7 +44,6 @@ Collection( id=uuid.UUID(int=3), name="test_collection_3", - topic=f"persistent://{PULSAR_TENANT}/{PULSAR_NAMESPACE}/chroma_log_14", metadata={"test_str": "str3", "test_int": 3, "test_float": 3.3}, dimension=None, database=DEFAULT_DATABASE, @@ -62,21 +52,12 @@ ] -class MockAssignmentPolicy(Component): - def assign_collection(self, collection_id: uuid.UUID) -> str: - for collection in sample_collections: - if collection["id"] == collection_id: - return collection["topic"] - raise ValueError(f"Unknown collection ID: {collection_id}") - - def sqlite() -> Generator[SysDB, None, None]: """Fixture generator for sqlite DB""" db = SqliteDB( System( Settings( allow_reset=True, - chroma_collection_assignment_policy_impl="chromadb.test.db.test_system.MockAssignmentPolicy", ) ) ) @@ -94,7 +75,6 @@ def sqlite_persistent() -> Generator[SysDB, None, None]: allow_reset=True, is_persistent=True, persist_directory=save_path, - chroma_collection_assignment_policy_impl="chromadb.test.db.test_system.MockAssignmentPolicy", ) ) ) @@ -111,7 +91,6 @@ def grpc_with_mock_server() -> Generator[SysDB, None, None]: system = System( Settings( allow_reset=True, - chroma_collection_assignment_policy_impl="chromadb.test.db.test_system.MockAssignmentPolicy", chroma_server_grpc_port=50051, ) ) @@ -120,13 +99,14 @@ def grpc_with_mock_server() -> Generator[SysDB, None, None]: system.start() client.reset_and_wait_for_ready() yield client + system.stop() def grpc_with_real_server() -> Generator[SysDB, None, None]: system = System( Settings( allow_reset=True, - chroma_collection_assignment_policy_impl="chromadb.test.db.test_system.MockAssignmentPolicy", + chroma_server_grpc_port=50051, ) ) client = system.instance(GrpcSysDB) @@ -177,26 +157,11 @@ def test_create_get_delete_collections(sysdb: SysDB) -> None: result = sysdb.get_collections(name=collection["name"]) assert result == [collection] - # Find by topic - for collection in sample_collections: - result = sysdb.get_collections(topic=collection["topic"]) - assert collection in result - # Find by id for collection in sample_collections: result = sysdb.get_collections(id=collection["id"]) assert result == [collection] - # Find by id and topic (positive case) - for collection in sample_collections: - result = sysdb.get_collections(id=collection["id"], topic=collection["topic"]) - assert result == [collection] - - # find by id and topic (negative case) - for collection in sample_collections: - result = sysdb.get_collections(id=collection["id"], topic="other_topic") - assert result == [] - # Delete c1 = sample_collections[0] sysdb.delete_collection(c1["id"]) @@ -218,7 +183,6 @@ def test_update_collections(sysdb: SysDB) -> None: coll = Collection( name=sample_collections[0]["name"], id=sample_collections[0]["id"], - topic=sample_collections[0]["topic"], metadata=sample_collections[0]["metadata"], dimension=sample_collections[0]["dimension"], database=DEFAULT_DATABASE, @@ -240,12 +204,6 @@ def test_update_collections(sysdb: SysDB) -> None: result = sysdb.get_collections(name=coll["name"]) assert result == [coll] - # Update topic - coll["topic"] = "new_topic" - sysdb.update_collection(coll["id"], topic=coll["topic"]) - result = sysdb.get_collections(topic=coll["topic"]) - assert result == [coll] - # Update dimension coll["dimension"] = 128 sysdb.update_collection(coll["id"], dimension=coll["dimension"]) @@ -621,14 +579,12 @@ def test_get_database_with_tenants(sysdb: SysDB) -> None: id=uuid.UUID("00000000-d7d7-413b-92e1-731098a6e492"), type="test_type_a", scope=SegmentScope.VECTOR, - topic=None, collection=sample_collections[0]["id"], metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3}, ), Segment( id=uuid.UUID("11111111-d7d7-413b-92e1-731098a6e492"), type="test_type_b", - topic="test_topic_2", scope=SegmentScope.VECTOR, collection=sample_collections[1]["id"], metadata={"test_str": "str2", "test_int": 2, "test_float": 2.3}, @@ -636,7 +592,6 @@ def test_get_database_with_tenants(sysdb: SysDB) -> None: Segment( id=uuid.UUID("22222222-d7d7-413b-92e1-731098a6e492"), type="test_type_b", - topic="test_topic_3", scope=SegmentScope.METADATA, collection=None, metadata={"test_str": "str3", "test_int": 3, "test_float": 3.3}, @@ -719,7 +674,6 @@ def test_update_segment(sysdb: SysDB) -> None: id=uuid.uuid4(), type="test_type_a", scope=SegmentScope.VECTOR, - topic="test_topic_a", collection=sample_collections[0]["id"], metadata=metadata, ) @@ -732,52 +686,56 @@ def test_update_segment(sysdb: SysDB) -> None: sysdb.create_segment(segment) - # Update topic to new value - segment["topic"] = "new_topic" - sysdb.update_segment(segment["id"], topic=segment["topic"]) + # TODO: revisit update segment - push collection id + result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] - # Update topic to None - segment["topic"] = None - sysdb.update_segment(segment["id"], topic=segment["topic"]) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] # Update collection to new value segment["collection"] = sample_collections[1]["id"] sysdb.update_segment(segment["id"], collection=segment["collection"]) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] # Update collection to None segment["collection"] = None sysdb.update_segment(segment["id"], collection=segment["collection"]) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] # Add a new metadata key metadata["test_str2"] = "str2" sysdb.update_segment(segment["id"], metadata={"test_str2": "str2"}) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] # Update a metadata key metadata["test_str"] = "str3" sysdb.update_segment(segment["id"], metadata={"test_str": "str3"}) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] # Delete a metadata key del metadata["test_str"] sysdb.update_segment(segment["id"], metadata={"test_str": None}) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] # Delete all metadata keys segment["metadata"] = None sysdb.update_segment(segment["id"], metadata=None) result = sysdb.get_segments(id=segment["id"]) + result[0]["collection"] = segment["collection"] assert result == [segment] diff --git a/chromadb/test/ef/test_ollama_ef.py b/chromadb/test/ef/test_ollama_ef.py new file mode 100644 index 00000000000..d44f1e8e6d1 --- /dev/null +++ b/chromadb/test/ef/test_ollama_ef.py @@ -0,0 +1,34 @@ +import os + +import pytest +import requests +from requests import HTTPError +from requests.exceptions import ConnectionError + +from chromadb.utils.embedding_functions import OllamaEmbeddingFunction + + +def test_ollama() -> None: + """ + To set up the Ollama server, follow instructions at: https://github.com/ollama/ollama?tab=readme-ov-file + Export the OLLAMA_SERVER_URL and OLLAMA_MODEL environment variables. + """ + if ( + os.environ.get("OLLAMA_SERVER_URL") is None + or os.environ.get("OLLAMA_MODEL") is None + ): + pytest.skip( + "OLLAMA_SERVER_URL or OLLAMA_MODEL environment variable not set. Skipping test." + ) + try: + response = requests.get(os.environ.get("OLLAMA_SERVER_URL", "")) + # If the response was successful, no Exception will be raised + response.raise_for_status() + except (HTTPError, ConnectionError): + pytest.skip("Ollama server not running. Skipping test.") + ef = OllamaEmbeddingFunction( + model_name=os.environ.get("OLLAMA_MODEL") or "nomic-embed-text", + url=f"{os.environ.get('OLLAMA_SERVER_URL')}/embeddings", + ) + embeddings = ef(["Here is an article about llamas...", "this is another article"]) + assert len(embeddings) == 2 diff --git a/chromadb/test/ingest/test_producer_consumer.py b/chromadb/test/ingest/test_producer_consumer.py index 199afde60de..d44faca5894 100644 --- a/chromadb/test/ingest/test_producer_consumer.py +++ b/chromadb/test/ingest/test_producer_consumer.py @@ -2,6 +2,7 @@ import os import shutil import tempfile +from uuid import UUID import pytest from itertools import count from typing import ( @@ -17,18 +18,16 @@ ) from chromadb.ingest import Producer, Consumer from chromadb.db.impl.sqlite import SqliteDB -from chromadb.ingest.impl.utils import create_topic_name from chromadb.test.conftest import ProducerFn from chromadb.types import ( - SubmitEmbeddingRecord, + OperationRecord, Operation, - EmbeddingRecord, + LogRecord, ScalarEncoding, ) from chromadb.config import System, Settings from pytest import FixtureRequest, approx from asyncio import Event, wait_for, TimeoutError -import uuid def sqlite() -> Generator[Tuple[Producer, Consumer], None, None]: @@ -54,29 +53,11 @@ def sqlite_persistent() -> Generator[Tuple[Producer, Consumer], None, None]: shutil.rmtree(save_path) -def pulsar() -> Generator[Tuple[Producer, Consumer], None, None]: - """Fixture generator for pulsar Producer + Consumer. This fixture requires a running - pulsar cluster. You can use bin/cluster-test.sh to start a standalone pulsar and run this test. - Assumes pulsar_broker_url etc is set from the environment variables like PULSAR_BROKER_URL. - """ - system = System( - Settings( - allow_reset=True, - chroma_producer_impl="chromadb.ingest.impl.pulsar.PulsarProducer", - chroma_consumer_impl="chromadb.ingest.impl.pulsar.PulsarConsumer", - ) - ) - producer = system.require(Producer) - consumer = system.require(Consumer) - system.start() - yield producer, consumer - system.stop() - - def fixtures() -> List[Callable[[], Generator[Tuple[Producer, Consumer], None, None]]]: fixtures = [sqlite, sqlite_persistent] if "CHROMA_CLUSTER_TEST_ONLY" in os.environ: - fixtures = [pulsar] + # TODO: We should add the new log service here + fixtures = [] return fixtures @@ -89,8 +70,8 @@ def producer_consumer( @pytest.fixture(scope="module") -def sample_embeddings() -> Iterator[SubmitEmbeddingRecord]: - def create_record(i: int) -> SubmitEmbeddingRecord: +def sample_embeddings() -> Iterator[OperationRecord]: + def create_record(i: int) -> OperationRecord: vector = [i + i * 0.1, i + 1 + i * 0.1] metadata: Optional[Dict[str, Union[str, int, float]]] if i % 2 == 0: @@ -98,13 +79,12 @@ def create_record(i: int) -> SubmitEmbeddingRecord: else: metadata = {"str_key": f"value_{i}", "int_key": i, "float_key": i + i * 0.1} - record = SubmitEmbeddingRecord( + record = OperationRecord( id=f"embedding_{i}", embedding=vector, encoding=ScalarEncoding.FLOAT32, metadata=metadata, operation=Operation.ADD, - collection_id=uuid.uuid4(), ) return record @@ -112,7 +92,7 @@ def create_record(i: int) -> SubmitEmbeddingRecord: class CapturingConsumeFn: - embeddings: List[EmbeddingRecord] + embeddings: List[LogRecord] waiters: List[Tuple[int, Event]] def __init__(self) -> None: @@ -124,14 +104,14 @@ def __init__(self) -> None: self.waiters = [] self._loop = asyncio.get_event_loop() - def __call__(self, embeddings: Sequence[EmbeddingRecord]) -> None: + def __call__(self, embeddings: Sequence[LogRecord]) -> None: self.embeddings.extend(embeddings) for n, event in self.waiters: if len(self.embeddings) >= n: # event.set() is not thread safe, so we need to call it in the main event loop self._loop.call_soon_threadsafe(event.set) - async def get(self, n: int, timeout_secs: int = 10) -> Sequence[EmbeddingRecord]: + async def get(self, n: int, timeout_secs: int = 10) -> Sequence[LogRecord]: "Wait until at least N embeddings are available, then return all embeddings" if len(self.embeddings) >= n: return self.embeddings[:n] @@ -149,41 +129,38 @@ def assert_approx_equal(a: Sequence[float], b: Sequence[float]) -> None: def assert_records_match( - inserted_records: Sequence[SubmitEmbeddingRecord], - consumed_records: Sequence[EmbeddingRecord], + inserted_records: Sequence[OperationRecord], + consumed_records: Sequence[LogRecord], ) -> None: """Given a list of inserted and consumed records, make sure they match""" assert len(consumed_records) == len(inserted_records) for inserted, consumed in zip(inserted_records, consumed_records): - assert inserted["id"] == consumed["id"] - assert inserted["operation"] == consumed["operation"] - assert inserted["encoding"] == consumed["encoding"] - assert inserted["metadata"] == consumed["metadata"] + assert inserted["id"] == consumed["operation_record"]["id"] + assert inserted["operation"] == consumed["operation_record"]["operation"] + assert inserted["encoding"] == consumed["operation_record"]["encoding"] + assert inserted["metadata"] == consumed["operation_record"]["metadata"] if inserted["embedding"] is not None: - assert consumed["embedding"] is not None - assert_approx_equal(inserted["embedding"], consumed["embedding"]) - - -def full_topic_name(topic_name: str) -> str: - return create_topic_name("default", "default", topic_name) + assert consumed["operation_record"]["embedding"] is not None + assert_approx_equal( + inserted["embedding"], consumed["operation_record"]["embedding"] + ) @pytest.mark.asyncio async def test_backfill( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name = full_topic_name("test_topic") - producer.create_topic(topic_name) - embeddings = produce_fns(producer, topic_name, sample_embeddings, 3)[0] + collection_id = UUID("00000000-0000-0000-0000-000000000000") + embeddings = produce_fns(producer, collection_id, sample_embeddings, 3)[0] consume_fn = CapturingConsumeFn() - consumer.subscribe(topic_name, consume_fn, start=consumer.min_seqid()) + consumer.subscribe(collection_id, consume_fn, start=consumer.min_seqid()) recieved = await consume_fn.get(3) assert_records_match(embeddings, recieved) @@ -192,61 +169,57 @@ async def test_backfill( @pytest.mark.asyncio async def test_notifications( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name = full_topic_name("test_topic") - - producer.create_topic(topic_name) + collection_id = UUID("00000000-0000-0000-0000-000000000000") - embeddings: List[SubmitEmbeddingRecord] = [] + embeddings: List[OperationRecord] = [] consume_fn = CapturingConsumeFn() - consumer.subscribe(topic_name, consume_fn, start=consumer.min_seqid()) + consumer.subscribe(collection_id, consume_fn, start=consumer.min_seqid()) for i in range(10): e = next(sample_embeddings) embeddings.append(e) - producer.submit_embedding(topic_name, e) + producer.submit_embedding(collection_id, e) received = await consume_fn.get(i + 1) assert_records_match(embeddings, received) @pytest.mark.asyncio -async def test_multiple_topics( +async def test_multiple_collections( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name_1 = full_topic_name("test_topic_1") - topic_name_2 = full_topic_name("test_topic_2") - producer.create_topic(topic_name_1) - producer.create_topic(topic_name_2) + collection_1 = UUID("00000000-0000-0000-0000-000000000001") + collection_2 = UUID("00000000-0000-0000-0000-000000000002") - embeddings_1: List[SubmitEmbeddingRecord] = [] - embeddings_2: List[SubmitEmbeddingRecord] = [] + embeddings_1: List[OperationRecord] = [] + embeddings_2: List[OperationRecord] = [] consume_fn_1 = CapturingConsumeFn() consume_fn_2 = CapturingConsumeFn() - consumer.subscribe(topic_name_1, consume_fn_1, start=consumer.min_seqid()) - consumer.subscribe(topic_name_2, consume_fn_2, start=consumer.min_seqid()) + consumer.subscribe(collection_1, consume_fn_1, start=consumer.min_seqid()) + consumer.subscribe(collection_2, consume_fn_2, start=consumer.min_seqid()) for i in range(10): e_1 = next(sample_embeddings) embeddings_1.append(e_1) - producer.submit_embedding(topic_name_1, e_1) + producer.submit_embedding(collection_1, e_1) results_2 = await consume_fn_1.get(i + 1) assert_records_match(embeddings_1, results_2) e_2 = next(sample_embeddings) embeddings_2.append(e_2) - producer.submit_embedding(topic_name_2, e_2) + producer.submit_embedding(collection_2, e_2) results_2 = await consume_fn_2.get(i + 1) assert_records_match(embeddings_2, results_2) @@ -254,28 +227,27 @@ async def test_multiple_topics( @pytest.mark.asyncio async def test_start_seq_id( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name = full_topic_name("test_topic") - producer.create_topic(topic_name) + collection = UUID("00000000-0000-0000-0000-000000000000") consume_fn_1 = CapturingConsumeFn() consume_fn_2 = CapturingConsumeFn() - consumer.subscribe(topic_name, consume_fn_1, start=consumer.min_seqid()) + consumer.subscribe(collection, consume_fn_1, start=consumer.min_seqid()) - embeddings = produce_fns(producer, topic_name, sample_embeddings, 5)[0] + embeddings = produce_fns(producer, collection, sample_embeddings, 5)[0] results_1 = await consume_fn_1.get(5) assert_records_match(embeddings, results_1) - start = consume_fn_1.embeddings[-1]["seq_id"] - consumer.subscribe(topic_name, consume_fn_2, start=start) - second_embeddings = produce_fns(producer, topic_name, sample_embeddings, 5)[0] + start = consume_fn_1.embeddings[-1]["log_offset"] + consumer.subscribe(collection, consume_fn_2, start=start) + second_embeddings = produce_fns(producer, collection, sample_embeddings, 5)[0] assert isinstance(embeddings, list) embeddings.extend(second_embeddings) results_2 = await consume_fn_2.get(5) @@ -285,27 +257,26 @@ async def test_start_seq_id( @pytest.mark.asyncio async def test_end_seq_id( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name = full_topic_name("test_topic") - producer.create_topic(topic_name) + collection = UUID("00000000-0000-0000-0000-000000000000") consume_fn_1 = CapturingConsumeFn() consume_fn_2 = CapturingConsumeFn() - consumer.subscribe(topic_name, consume_fn_1, start=consumer.min_seqid()) + consumer.subscribe(collection, consume_fn_1, start=consumer.min_seqid()) - embeddings = produce_fns(producer, topic_name, sample_embeddings, 10)[0] + embeddings = produce_fns(producer, collection, sample_embeddings, 10)[0] results_1 = await consume_fn_1.get(10) assert_records_match(embeddings, results_1) - end = consume_fn_1.embeddings[-5]["seq_id"] - consumer.subscribe(topic_name, consume_fn_2, start=consumer.min_seqid(), end=end) + end = consume_fn_1.embeddings[-5]["log_offset"] + consumer.subscribe(collection, consume_fn_2, start=consumer.min_seqid(), end=end) results_2 = await consume_fn_2.get(6) assert_records_match(embeddings[:6], results_2) @@ -318,29 +289,28 @@ async def test_end_seq_id( @pytest.mark.asyncio async def test_submit_batch( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name = full_topic_name("test_topic") + collection = UUID("00000000-0000-0000-0000-000000000000") embeddings = [next(sample_embeddings) for _ in range(100)] - producer.create_topic(topic_name) - producer.submit_embeddings(topic_name, embeddings=embeddings) + producer.submit_embeddings(collection, embeddings=embeddings) consume_fn = CapturingConsumeFn() - consumer.subscribe(topic_name, consume_fn, start=consumer.min_seqid()) + consumer.subscribe(collection, consume_fn, start=consumer.min_seqid()) recieved = await consume_fn.get(100) assert_records_match(embeddings, recieved) @pytest.mark.asyncio -async def test_multiple_topics_batch( +async def test_multiple_collections_batch( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer, consumer = producer_consumer @@ -350,14 +320,13 @@ async def test_multiple_topics_batch( N_TOPICS = 2 consume_fns = [CapturingConsumeFn() for _ in range(N_TOPICS)] for i in range(N_TOPICS): - producer.create_topic(full_topic_name(f"test_topic_{i}")) consumer.subscribe( - full_topic_name(f"test_topic_{i}"), + UUID(f"00000000-0000-0000-0000-00000000000{i}"), consume_fns[i], start=consumer.min_seqid(), ) - embeddings_n: List[List[SubmitEmbeddingRecord]] = [[] for _ in range(N_TOPICS)] + embeddings_n: List[List[OperationRecord]] = [[] for _ in range(N_TOPICS)] PRODUCE_BATCH_SIZE = 10 N_TO_PRODUCE = 100 @@ -367,7 +336,7 @@ async def test_multiple_topics_batch( embeddings_n[n].extend( produce_fns( producer, - full_topic_name(f"test_topic_{n}"), + UUID(f"00000000-0000-0000-0000-00000000000{n}"), sample_embeddings, PRODUCE_BATCH_SIZE, )[0] @@ -380,25 +349,25 @@ async def test_multiple_topics_batch( @pytest.mark.asyncio async def test_max_batch_size( producer_consumer: Tuple[Producer, Consumer], - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], ) -> None: producer, consumer = producer_consumer producer.reset_state() consumer.reset_state() - topic_name = full_topic_name("test_topic") + collection = UUID("00000000-0000-0000-0000-000000000000") max_batch_size = producer.max_batch_size assert max_batch_size > 0 # Make sure that we can produce a batch of size max_batch_size embeddings = [next(sample_embeddings) for _ in range(max_batch_size)] consume_fn = CapturingConsumeFn() - consumer.subscribe(topic_name, consume_fn, start=consumer.min_seqid()) - producer.submit_embeddings(topic_name, embeddings=embeddings) + consumer.subscribe(collection, consume_fn, start=consumer.min_seqid()) + producer.submit_embeddings(collection, embeddings=embeddings) received = await consume_fn.get(max_batch_size, timeout_secs=120) assert_records_match(embeddings, received) embeddings = [next(sample_embeddings) for _ in range(max_batch_size + 1)] # Make sure that we can't produce a batch of size > max_batch_size with pytest.raises(ValueError) as e: - producer.submit_embeddings(topic_name, embeddings=embeddings) + producer.submit_embeddings(collection, embeddings=embeddings) assert "Cannot submit more than" in str(e.value) diff --git a/chromadb/test/openssl.cnf b/chromadb/test/openssl.cnf new file mode 100644 index 00000000000..11704076bd4 --- /dev/null +++ b/chromadb/test/openssl.cnf @@ -0,0 +1,12 @@ +[req] +distinguished_name = req_distinguished_name +x509_extensions = usr_cert + +[req_distinguished_name] +CN = localhost + +[usr_cert] +subjectAltName = @alt_names + +[alt_names] +DNS.1 = localhost \ No newline at end of file diff --git a/chromadb/test/property/strategies.py b/chromadb/test/property/strategies.py index 5a4c0d905cc..79f7a1fd326 100644 --- a/chromadb/test/property/strategies.py +++ b/chromadb/test/property/strategies.py @@ -3,6 +3,7 @@ import hypothesis.strategies as st from typing import Any, Optional, List, Dict, Union, cast from typing_extensions import TypedDict +import uuid import numpy as np import numpy.typing as npt import chromadb.api.types as types @@ -235,16 +236,37 @@ def embedding_function_strategy( @dataclass -class Collection: +class ExternalCollection: + """ + An external view of a collection. + + This strategy only contains information about a collection that a client of Chroma + sees -- that is, it contains none of Chroma's internal bookkeeping. It should + be used to test the API and client code. + """ + name: str metadata: Optional[types.Metadata] + embedding_function: Optional[types.EmbeddingFunction[Embeddable]] + + +@dataclass +class Collection(ExternalCollection): + """ + An internal view of a collection. + + This strategy contains all the information Chroma uses internally to manage a + collection. It is a superset of ExternalCollection and should be used to test + internal Chroma logic. + """ + + id: uuid.UUID dimension: int dtype: npt.DTypeLike known_metadata_keys: types.Metadata known_document_keywords: List[str] has_documents: bool = False has_embeddings: bool = False - embedding_function: Optional[types.EmbeddingFunction[Embeddable]] = None @st.composite @@ -309,6 +331,7 @@ def collections( embedding_function = draw(embedding_function_strategy(dimension, dtype)) return Collection( + id=uuid.uuid4(), name=name, metadata=metadata, dimension=dimension, diff --git a/chromadb/test/property/test_collections.py b/chromadb/test/property/test_collections.py index 844476aa8ea..3aa2e454c07 100644 --- a/chromadb/test/property/test_collections.py +++ b/chromadb/test/property/test_collections.py @@ -16,9 +16,13 @@ ) from typing import Dict, Optional +import numpy +import uuid +from chromadb.test.property.strategies import hashing_embedding_function + class CollectionStateMachine(RuleBasedStateMachine): - collections: Bundle[strategies.Collection] + collections: Bundle[strategies.ExternalCollection] _model: Dict[str, Optional[types.CollectionMetadata]] collections = Bundle("collections") @@ -35,8 +39,8 @@ def initialize(self) -> None: @rule(target=collections, coll=strategies.collections()) def create_coll( - self, coll: strategies.Collection - ) -> MultipleResults[strategies.Collection]: + self, coll: strategies.ExternalCollection + ) -> MultipleResults[strategies.ExternalCollection]: # Metadata can either be None or a non-empty dict if coll.name in self.model or ( coll.metadata is not None and len(coll.metadata) == 0 @@ -61,7 +65,7 @@ def create_coll( return multiple(coll) @rule(coll=collections) - def get_coll(self, coll: strategies.Collection) -> None: + def get_coll(self, coll: strategies.ExternalCollection) -> None: if coll.name in self.model: c = self.api.get_collection(name=coll.name) assert c.name == coll.name @@ -71,7 +75,7 @@ def get_coll(self, coll: strategies.Collection) -> None: self.api.get_collection(name=coll.name) @rule(coll=consumes(collections)) - def delete_coll(self, coll: strategies.Collection) -> None: + def delete_coll(self, coll: strategies.ExternalCollection) -> None: if coll.name in self.model: self.api.delete_collection(name=coll.name) self.delete_from_model(coll.name) @@ -119,9 +123,9 @@ def list_collections_with_limit_offset(self, limit: int, offset: int) -> None: ) def get_or_create_coll( self, - coll: strategies.Collection, + coll: strategies.ExternalCollection, new_metadata: Optional[types.Metadata], - ) -> MultipleResults[strategies.Collection]: + ) -> MultipleResults[strategies.ExternalCollection]: # Cases for get_or_create # Case 0 @@ -185,17 +189,18 @@ def get_or_create_coll( ) def modify_coll( self, - coll: strategies.Collection, + coll: strategies.ExternalCollection, new_metadata: types.Metadata, new_name: Optional[str], - ) -> MultipleResults[strategies.Collection]: + ) -> MultipleResults[strategies.ExternalCollection]: if coll.name not in self.model: with pytest.raises(Exception): c = self.api.get_collection(name=coll.name) return multiple() c = self.api.get_collection(name=coll.name) - + _metadata: Optional[Mapping[str, Any]] = self.model[coll.name] + _name: str = coll.name if new_metadata is not None: if len(new_metadata) == 0: with pytest.raises(Exception): @@ -206,7 +211,7 @@ def modify_coll( ) return multiple() coll.metadata = new_metadata - self.set_model(coll.name, coll.metadata) + _metadata = new_metadata if new_name is not None: if new_name in self.model and new_name != coll.name: @@ -214,12 +219,12 @@ def modify_coll( c.modify(metadata=new_metadata, name=new_name) return multiple() - prev_metadata = self.model[coll.name] self.delete_from_model(coll.name) - self.set_model(new_name, prev_metadata) coll.name = new_name + _name = new_name - c.modify(metadata=new_metadata, name=new_name) + self.set_model(_name, _metadata) + c.modify(metadata=_metadata, name=_name) c = self.api.get_collection(name=coll.name) assert c.name == coll.name @@ -227,7 +232,9 @@ def modify_coll( return multiple(coll) def set_model( - self, name: str, metadata: Optional[types.CollectionMetadata] + self, + name: str, + metadata: Optional[types.CollectionMetadata], ) -> None: model = self.model model[name] = metadata @@ -244,3 +251,87 @@ def model(self) -> Dict[str, Optional[types.CollectionMetadata]]: def test_collections(caplog: pytest.LogCaptureFixture, api: ClientAPI) -> None: caplog.set_level(logging.ERROR) run_state_machine_as_test(lambda: CollectionStateMachine(api)) # type: ignore + + +# Below are tests that have failed in the past. If your test fails, please add +# it to protect against regressions in the test harness itself. If you need +# help doing so, talk to ben. + + +def test_previously_failing_one(api: ClientAPI) -> None: + state = CollectionStateMachine(api) + state.initialize() + # I don't know why the typechecker is red here. This code is correct and is + # pulled from the logs. + (v1,) = state.get_or_create_coll( + coll=strategies.ExternalCollection( + name='jjn2yjLW1zp2T\n', + metadata=None, + embedding_function=hashing_embedding_function(dtype=numpy.float32, dim=863) + ), + new_metadata=None + ) + (v6,) = state.get_or_create_coll( + coll=strategies.ExternalCollection( + name='jjn2yjLW1zp2T\n', + metadata=None, + embedding_function=hashing_embedding_function(dtype=numpy.float32, dim=863) + ), + new_metadata=None + ) + state.modify_coll( + coll=v1, + new_metadata={'7': -1281, 'fGe': -0.0, 'K5j': 'im'}, + new_name=None + ) + state.modify_coll( + coll=v6, + new_metadata=None, + new_name=None + ) + + +# https://github.com/chroma-core/chroma/commit/cf476d70f0cebb7c87cb30c7172ba74d6ea175cd#diff-e81868b665d149bb315d86890dea6fc6a9fc9fc9ea3089aa7728142b54f622c5R210 +def test_previously_failing_two(api: ClientAPI) -> None: + state = CollectionStateMachine(api) + state.initialize() + (v13,) = state.get_or_create_coll( + coll=strategies.ExternalCollection( + name='C1030', + metadata={}, + embedding_function=hashing_embedding_function(dim=2, dtype=numpy.float32) + ), + new_metadata=None + ) + (v15,) = state.modify_coll( + coll=v13, + new_metadata={'0': '10', '40': '0', 'p1nviWeL7fO': 'qN', '7b': 'YS', 'VYWq4LEMWjCo': True}, + new_name='OF5F0MzbQg\n' + ) + state.get_or_create_coll( + coll=strategies.ExternalCollection( + name='VS0QGh', + metadata={'h': 5.681951615025145e-227, 'A1': 61126, 'uhUhLEEMfeC_kN': 2147483647, 'weF': 'pSP', 'B3DSaP': False, '6H533K': 1.192092896e-07}, + embedding_function=hashing_embedding_function(dim=1915, dtype=numpy.float32) + ), + new_metadata={'xVW09xUpDZA': 31734, + 'g': 1.1, + 'n1dUTalF-MY': -1000000.0, + 'y': 'G3EtXTZ', + 'ugXZ_hK': 5494 + } + ) + (v17,) = state.modify_coll( + coll=v15, + new_metadata={'L35J2S': 'K0l026'}, + new_name='Ai1\n' + ) + (v18,) = state.get_or_create_coll(coll=v13, new_metadata=None) + state.get_or_create_coll( + coll=strategies.ExternalCollection( + name='VS0QGh', + metadata=None, + embedding_function=hashing_embedding_function(dim=326, dtype=numpy.float16) + ), + new_metadata=None + ) \ No newline at end of file diff --git a/chromadb/test/property/test_embeddings.py b/chromadb/test/property/test_embeddings.py index cfb2c93fa52..bf3e882184f 100644 --- a/chromadb/test/property/test_embeddings.py +++ b/chromadb/test/property/test_embeddings.py @@ -455,3 +455,10 @@ def test_autocasting_validate_embeddings_incompatible_types( validate_embeddings(Collection._normalize_embeddings(embds)) assert "Expected each value in the embedding to be a int or float" in str(e) + + +def test_0dim_embedding_validation() -> None: + embds = [[]] + with pytest.raises(ValueError) as e: + validate_embeddings(embds) + assert "Expected each embedding in the embeddings to be a non-empty list" in str(e) \ No newline at end of file diff --git a/chromadb/test/property/test_segment_manager.py b/chromadb/test/property/test_segment_manager.py new file mode 100644 index 00000000000..ff5e057dff4 --- /dev/null +++ b/chromadb/test/property/test_segment_manager.py @@ -0,0 +1,128 @@ +import uuid + +import pytest +import chromadb.test.property.strategies as strategies +from unittest.mock import patch +from dataclasses import asdict +import random +from hypothesis.stateful import ( + Bundle, + RuleBasedStateMachine, + rule, + initialize, + multiple, + precondition, + invariant, + run_state_machine_as_test, + MultipleResults, +) +from typing import Dict +from chromadb.segment import ( + VectorReader +) +from chromadb.segment import SegmentManager + +from chromadb.segment.impl.manager.local import LocalSegmentManager +from chromadb.types import SegmentScope +from chromadb.db.system import SysDB +from chromadb.config import System, get_class + +# Memory limit use for testing +memory_limit = 100 + +# Helper class to keep tract of the last use id +class LastUse: + def __init__(self, n: int): + self.n = n + self.store = [] + + def add(self, id: uuid.UUID): + if id in self.store: + self.store.remove(id) + self.store.append(id) + else: + self.store.append(id) + while len(self.store) > self.n: + self.store.pop(0) + return self.store + + def reset(self): + self.store = [] + + +class SegmentManagerStateMachine(RuleBasedStateMachine): + collections: Bundle[strategies.Collection] + collections = Bundle("collections") + collection_size_store: Dict[uuid.UUID, int] = {} + segment_collection: Dict[uuid.UUID, uuid.UUID] = {} + + def __init__(self, system: System): + super().__init__() + self.segment_manager = system.require(SegmentManager) + self.segment_manager.start() + self.segment_manager.reset_state() + self.last_use = LastUse(n=40) + self.collection_created_counter = 0 + self.sysdb = system.require(SysDB) + self.system = system + + @invariant() + def last_queried_segments_should_be_in_cache(self): + cache_sum = 0 + index = 0 + for id in reversed(self.last_use.store): + cache_sum += self.collection_size_store[id] + if cache_sum >= memory_limit and index is not 0: + break + assert id in self.segment_manager.segment_cache[SegmentScope.VECTOR].cache + index += 1 + + @invariant() + @precondition(lambda self: self.system.settings.is_persistent is True) + def cache_should_not_be_bigger_than_settings(self): + segment_sizes = {id: self.collection_size_store[id] for id in self.segment_manager.segment_cache[SegmentScope.VECTOR].cache} + total_size = sum(segment_sizes.values()) + if len(segment_sizes) != 1: + assert total_size <= memory_limit + + @initialize() + def initialize(self) -> None: + self.segment_manager.reset_state() + self.segment_manager.start() + self.collection_created_counter = 0 + self.last_use.reset() + + @rule(target=collections, coll=strategies.collections()) + @precondition(lambda self: self.collection_created_counter <= 50) + def create_segment( + self, coll: strategies.Collection + ) -> MultipleResults[strategies.Collection]: + segments = self.segment_manager.create_segments(asdict(coll)) + for segment in segments: + self.sysdb.create_segment(segment) + self.segment_collection[segment["id"]] = coll.id + self.collection_created_counter += 1 + self.collection_size_store[coll.id] = random.randint(0, memory_limit) + return multiple(coll) + + @rule(coll=collections) + def get_segment(self, coll: strategies.Collection) -> None: + segment = self.segment_manager.get_segment(collection_id=coll.id, type=VectorReader) + self.last_use.add(coll.id) + assert segment is not None + + + @staticmethod + def mock_directory_size(directory: str): + path_id = directory.split("/").pop() + collection_id = SegmentManagerStateMachine.segment_collection[uuid.UUID(path_id)] + return SegmentManagerStateMachine.collection_size_store[collection_id] + + +@patch('chromadb.segment.impl.manager.local.get_directory_size', SegmentManagerStateMachine.mock_directory_size) +def test_segment_manager(caplog: pytest.LogCaptureFixture, system: System) -> None: + system.settings.chroma_memory_limit_bytes = memory_limit + system.settings.chroma_segment_cache_policy = "LRU" + + run_state_machine_as_test( + lambda: SegmentManagerStateMachine(system=system)) diff --git a/chromadb/test/quota/test_static_quota_enforcer.py b/chromadb/test/quota/test_static_quota_enforcer.py new file mode 100644 index 00000000000..245e9ba2e80 --- /dev/null +++ b/chromadb/test/quota/test_static_quota_enforcer.py @@ -0,0 +1,78 @@ +import random +import string +from typing import Optional, List, Tuple, Any +from unittest.mock import patch + +from chromadb.config import System, Settings +from chromadb.quota import QuotaEnforcer, Resource +import pytest + + +def generate_random_string(size: int) -> str: + return ''.join(random.choices(string.ascii_letters + string.digits, k=size)) + +def mock_get_for_subject(self, resource: Resource, subject: Optional[str] = "", tier: Optional[str] = "") -> Optional[ + int]: + """Mock function to simulate quota retrieval.""" + return 10 + + +def run_static_checks(enforcer: QuotaEnforcer, test_cases: List[Tuple[Any, Optional[str]]], data_key: str): + """Generalized function to run static checks on different types of data.""" + for test_case in test_cases: + data, expected_error = test_case if len(test_case) == 2 else (test_case[0], None) + args = {data_key: [data]} + if expected_error: + with pytest.raises(Exception) as exc_info: + enforcer.static_check(**args) + assert expected_error in str(exc_info.value.resource) + else: + enforcer.static_check(**args) + + + +@pytest.fixture(scope="module") +def enforcer() -> QuotaEnforcer: + settings = Settings( + chroma_quota_provider_impl = "chromadb.quota.test_provider.QuotaProviderForTest" + ) + system = System(settings) + return system.require(QuotaEnforcer) + +@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject) +def test_static_enforcer_metadata(enforcer): + test_cases = [ + ({generate_random_string(20): generate_random_string(5)}, "METADATA_KEY_LENGTH"), + ({generate_random_string(5): generate_random_string(5)}, None), + ({generate_random_string(5): generate_random_string(20)}, "METADATA_VALUE_LENGTH"), + ({generate_random_string(5): generate_random_string(5)}, None) + ] + run_static_checks(enforcer, test_cases, 'metadatas') + + +@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject) +def test_static_enforcer_documents(enforcer): + test_cases = [ + (generate_random_string(20), "DOCUMENT_SIZE"), + (generate_random_string(5), None) + ] + run_static_checks(enforcer, test_cases, 'documents') + +@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject) +def test_static_enforcer_embeddings(enforcer): + test_cases = [ + (random.sample(range(1, 101), 100), "EMBEDDINGS_DIMENSION"), + (random.sample(range(1, 101), 5), None) + ] + run_static_checks(enforcer, test_cases, 'embeddings') + +# Should not raise an error if no quota provider is present +def test_enforcer_without_quota_provider(): + test_cases = [ + (random.sample(range(1, 101), 1), None), + (random.sample(range(1, 101), 5), None) + ] + settings = Settings() + system = System(settings) + enforcer = system.require(QuotaEnforcer) + run_static_checks(enforcer, test_cases, 'embeddings') diff --git a/chromadb/test/rate_limiting/test_rate_limiting.py b/chromadb/test/rate_limiting/test_rate_limiting.py new file mode 100644 index 00000000000..6f8fb7b8c42 --- /dev/null +++ b/chromadb/test/rate_limiting/test_rate_limiting.py @@ -0,0 +1,50 @@ +from typing import Optional +from unittest.mock import patch + +from chromadb.config import System, Settings, Component +from chromadb.quota import QuotaEnforcer, Resource +import pytest + +from chromadb.rate_limiting import rate_limit + + +class RateLimitingGym(Component): + def __init__(self, system: System): + super().__init__(system) + self.system = system + + @rate_limit(subject="bar", resource=Resource.DOCUMENT_SIZE) + def bench(self, foo: str, bar: str) -> str: + return foo + +def mock_get_for_subject(self, resource: Resource, subject: Optional[str] = "", tier: Optional[str] = "") -> Optional[ + int]: + """Mock function to simulate quota retrieval.""" + return 10 + +@pytest.fixture(scope="module") +def rate_limiting_gym() -> QuotaEnforcer: + settings = Settings( + chroma_quota_provider_impl="chromadb.quota.test_provider.QuotaProviderForTest", + chroma_rate_limiting_provider_impl="chromadb.rate_limiting.test_provider.RateLimitingTestProvider" + ) + system = System(settings) + return RateLimitingGym(system) + + +@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject) +@patch('chromadb.rate_limiting.test_provider.RateLimitingTestProvider.is_allowed', lambda self, key, quota, point=1: False) +def test_rate_limiting_should_raise(rate_limiting_gym: RateLimitingGym): + with pytest.raises(Exception) as exc_info: + rate_limiting_gym.bench("foo", "bar") + assert Resource.DOCUMENT_SIZE.value in str(exc_info.value.resource) + +@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject) +@patch('chromadb.rate_limiting.test_provider.RateLimitingTestProvider.is_allowed', lambda self, key, quota, point=1: True) +def test_rate_limiting_should_not_raise(rate_limiting_gym: RateLimitingGym): + assert rate_limiting_gym.bench(foo="foo", bar="bar") is "foo" + +@patch('chromadb.quota.test_provider.QuotaProviderForTest.get_for_subject', mock_get_for_subject) +@patch('chromadb.rate_limiting.test_provider.RateLimitingTestProvider.is_allowed', lambda self, key, quota, point=1: True) +def test_rate_limiting_should_not_raise(rate_limiting_gym: RateLimitingGym): + assert rate_limiting_gym.bench("foo", "bar") is "foo" \ No newline at end of file diff --git a/chromadb/test/segment/test_metadata.py b/chromadb/test/segment/test_metadata.py index ef6400b210e..3497d709dd8 100644 --- a/chromadb/test/segment/test_metadata.py +++ b/chromadb/test/segment/test_metadata.py @@ -2,13 +2,25 @@ import shutil import tempfile import pytest -from typing import Generator, List, Callable, Iterator, Dict, Optional, Union, Sequence +from typing import ( + Generator, + List, + Callable, + Iterator, + Dict, + Optional, + Union, + Sequence, + cast, +) + +from chromadb.api.types import validate_metadata from chromadb.config import System, Settings from chromadb.db.base import ParameterValue, get_sql from chromadb.db.impl.sqlite import SqliteDB from chromadb.test.conftest import ProducerFn from chromadb.types import ( - SubmitEmbeddingRecord, + OperationRecord, MetadataEmbeddingRecord, Operation, ScalarEncoding, @@ -61,8 +73,8 @@ def system(request: FixtureRequest) -> Generator[System, None, None]: @pytest.fixture(scope="function") -def sample_embeddings() -> Iterator[SubmitEmbeddingRecord]: - def create_record(i: int) -> SubmitEmbeddingRecord: +def sample_embeddings() -> Iterator[OperationRecord]: + def create_record(i: int) -> OperationRecord: vector = [i + i * 0.1, i + 1 + i * 0.1] metadata: Optional[Dict[str, Union[str, int, float, bool]]] if i == 0: @@ -80,13 +92,12 @@ def create_record(i: int) -> SubmitEmbeddingRecord: metadata["bool_key"] = False metadata["chroma:document"] = _build_document(i) - record = SubmitEmbeddingRecord( + record = OperationRecord( id=f"embedding_{i}", embedding=vector, encoding=ScalarEncoding.FLOAT32, metadata=metadata, operation=Operation.ADD, - collection_id=uuid.UUID(int=0), ) return record @@ -116,8 +127,7 @@ def _build_document(i: int) -> str: id=uuid.uuid4(), type="test_type", scope=SegmentScope.METADATA, - topic="persistent://test/test/test_topic_1", - collection=None, + collection=uuid.UUID(int=0), metadata=None, ) @@ -125,8 +135,7 @@ def _build_document(i: int) -> str: id=uuid.uuid4(), type="test_type", scope=SegmentScope.METADATA, - topic="persistent://test/test/test_topic_2", - collection=None, + collection=uuid.UUID(int=1), metadata=None, ) @@ -143,15 +152,17 @@ def sync(segment: MetadataReader, seq_id: SeqId) -> None: def test_insert_and_count( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) - max_id = produce_fns(producer, topic, sample_embeddings, 3)[1][-1] + max_id = produce_fns(producer, collection_id, sample_embeddings, 3)[1][-1] segment = SqliteMetadataSegment(system, segment_definition) segment.start() @@ -161,7 +172,7 @@ def test_insert_and_count( assert segment.count() == 3 for i in range(3): - max_id = producer.submit_embedding(topic, next(sample_embeddings)) + max_id = producer.submit_embedding(collection_id, next(sample_embeddings)) sync(segment, max_id) @@ -169,7 +180,7 @@ def test_insert_and_count( def assert_equiv_records( - expected: Sequence[SubmitEmbeddingRecord], actual: Sequence[MetadataEmbeddingRecord] + expected: Sequence[OperationRecord], actual: Sequence[MetadataEmbeddingRecord] ) -> None: assert len(expected) == len(actual) sorted_expected = sorted(expected, key=lambda r: r["id"]) @@ -181,14 +192,16 @@ def assert_equiv_records( def test_get( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) - embeddings, seq_ids = produce_fns(producer, topic, sample_embeddings, 10) + embeddings, seq_ids = produce_fns(producer, collection_id, sample_embeddings, 10) segment = SqliteMetadataSegment(system, segment_definition) segment.start() @@ -204,7 +217,6 @@ def test_get( # Get all records results = segment.get_metadata() - assert seq_ids == [r["seq_id"] for r in results] assert_equiv_records(embeddings, results) # get by ID @@ -283,17 +295,19 @@ def test_get( def test_fulltext( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) segment = SqliteMetadataSegment(system, segment_definition) segment.start() - max_id = produce_fns(producer, topic, sample_embeddings, 100)[1][-1] + max_id = produce_fns(producer, collection_id, sample_embeddings, 100)[1][-1] sync(segment, max_id) @@ -378,17 +392,19 @@ def test_fulltext( def test_delete( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) segment = SqliteMetadataSegment(system, segment_definition) segment.start() - embeddings, seq_ids = produce_fns(producer, topic, sample_embeddings, 10) + embeddings, seq_ids = produce_fns(producer, collection_id, sample_embeddings, 10) max_id = seq_ids[-1] sync(segment, max_id) @@ -398,17 +414,16 @@ def test_delete( assert_equiv_records(embeddings[:1], results) # Delete by ID - delete_embedding = SubmitEmbeddingRecord( + delete_embedding = OperationRecord( id="embedding_0", embedding=None, encoding=None, metadata=None, operation=Operation.DELETE, - collection_id=uuid.UUID(int=0), ) - max_id = produce_fns(producer, topic, (delete_embedding for _ in range(1)), 1)[1][ - -1 - ] + max_id = produce_fns( + producer, collection_id, (delete_embedding for _ in range(1)), 1 + )[1][-1] sync(segment, max_id) @@ -416,43 +431,42 @@ def test_delete( assert segment.get_metadata(ids=["embedding_0"]) == [] # Delete is idempotent - max_id = produce_fns(producer, topic, (delete_embedding for _ in range(1)), 1)[1][ - -1 - ] + max_id = produce_fns( + producer, collection_id, (delete_embedding for _ in range(1)), 1 + )[1][-1] sync(segment, max_id) assert segment.count() == 9 assert segment.get_metadata(ids=["embedding_0"]) == [] # re-add - max_id = producer.submit_embedding(topic, embeddings[0]) + max_id = producer.submit_embedding(collection_id, embeddings[0]) sync(segment, max_id) assert segment.count() == 10 results = segment.get_metadata(ids=["embedding_0"]) -def test_update( - system: System, sample_embeddings: Iterator[SubmitEmbeddingRecord] -) -> None: +def test_update(system: System, sample_embeddings: Iterator[OperationRecord]) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) segment = SqliteMetadataSegment(system, segment_definition) segment.start() - _test_update(sample_embeddings, producer, segment, topic, Operation.UPDATE) + _test_update(sample_embeddings, producer, segment, collection_id, Operation.UPDATE) # Update nonexisting ID - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="no_such_id", metadata={"foo": "bar"}, embedding=None, encoding=None, operation=Operation.UPDATE, - collection_id=uuid.UUID(int=0), ) - max_id = producer.submit_embedding(topic, update_record) + max_id = producer.submit_embedding(collection_id, update_record) sync(segment, max_id) results = segment.get_metadata(ids=["no_such_id"]) assert len(results) == 0 @@ -461,30 +475,31 @@ def test_update( def test_upsert( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) segment = SqliteMetadataSegment(system, segment_definition) segment.start() - _test_update(sample_embeddings, producer, segment, topic, Operation.UPSERT) + _test_update(sample_embeddings, producer, segment, collection_id, Operation.UPSERT) # upsert previously nonexisting ID - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="no_such_id", metadata={"foo": "bar"}, embedding=None, encoding=None, operation=Operation.UPSERT, - collection_id=uuid.UUID(int=0), ) max_id = produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, embeddings=(update_record for _ in range(1)), n=1, )[1][-1] @@ -494,10 +509,10 @@ def test_upsert( def _test_update( - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], producer: Producer, segment: MetadataReader, - topic: str, + collection_id: uuid.UUID, op: Operation, ) -> None: """test code common between update and upsert paths""" @@ -506,7 +521,7 @@ def _test_update( max_id = 0 for e in embeddings: - max_id = producer.submit_embedding(topic, e) + max_id = producer.submit_embedding(collection_id, e) sync(segment, max_id) @@ -514,15 +529,14 @@ def _test_update( assert_equiv_records(embeddings[:1], results) # Update embedding with no metadata - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="embedding_0", metadata={"chroma:document": "foo bar"}, embedding=None, encoding=None, operation=op, - collection_id=uuid.UUID(int=0), ) - max_id = producer.submit_embedding(topic, update_record) + max_id = producer.submit_embedding(collection_id, update_record) sync(segment, max_id) results = segment.get_metadata(ids=["embedding_0"]) assert results[0]["metadata"] == {"chroma:document": "foo bar"} @@ -530,15 +544,14 @@ def _test_update( assert results[0]["metadata"] == {"chroma:document": "foo bar"} # Update and overrwrite key - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="embedding_0", metadata={"chroma:document": "biz buz"}, embedding=None, encoding=None, operation=op, - collection_id=uuid.UUID(int=0), ) - max_id = producer.submit_embedding(topic, update_record) + max_id = producer.submit_embedding(collection_id, update_record) sync(segment, max_id) results = segment.get_metadata(ids=["embedding_0"]) assert results[0]["metadata"] == {"chroma:document": "biz buz"} @@ -548,29 +561,27 @@ def _test_update( assert len(results) == 0 # Update and add key - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="embedding_0", metadata={"baz": 42}, embedding=None, encoding=None, operation=op, - collection_id=uuid.UUID(int=0), ) - max_id = producer.submit_embedding(topic, update_record) + max_id = producer.submit_embedding(collection_id, update_record) sync(segment, max_id) results = segment.get_metadata(ids=["embedding_0"]) assert results[0]["metadata"] == {"chroma:document": "biz buz", "baz": 42} # Update and delete key - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="embedding_0", metadata={"chroma:document": None}, embedding=None, encoding=None, operation=op, - collection_id=uuid.UUID(int=0), ) - max_id = producer.submit_embedding(topic, update_record) + max_id = producer.submit_embedding(collection_id, update_record) sync(segment, max_id) results = segment.get_metadata(ids=["embedding_0"]) assert results[0]["metadata"] == {"baz": 42} @@ -580,17 +591,17 @@ def _test_update( def test_limit( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) - max_id = produce_fns(producer, topic, sample_embeddings, 3)[1][-1] + collection_id = cast(uuid.UUID, segment_definition["collection"]) + max_id = produce_fns(producer, collection_id, sample_embeddings, 3)[1][-1] - topic2 = str(segment_definition2["topic"]) - max_id2 = produce_fns(producer, topic2, sample_embeddings, 3)[1][-1] + collection_id_2 = cast(uuid.UUID, segment_definition2["collection"]) + max_id2 = produce_fns(producer, collection_id_2, sample_embeddings, 3)[1][-1] segment = SqliteMetadataSegment(system, segment_definition) segment.start() @@ -604,7 +615,7 @@ def test_limit( assert segment.count() == 3 for i in range(3): - max_id = producer.submit_embedding(topic, next(sample_embeddings)) + max_id = producer.submit_embedding(collection_id, next(sample_embeddings)) sync(segment, max_id) @@ -624,17 +635,19 @@ def test_limit( def test_delete_segment( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) segment = SqliteMetadataSegment(system, segment_definition) segment.start() - embeddings, seq_ids = produce_fns(producer, topic, sample_embeddings, 10) + embeddings, seq_ids = produce_fns(producer, collection_id, sample_embeddings, 10) max_id = seq_ids[-1] sync(segment, max_id) @@ -657,3 +670,90 @@ def test_delete_segment( res = cur.execute(sql, params) # assert that the segment is gone assert len(res.fetchall()) == 0 + + fts_t = Table("embedding_fulltext_search") + q_fts = ( + _db.querybuilder() + .from_(fts_t) + .select() + .where( + fts_t.rowid.isin( + _db.querybuilder() + .from_(t) + .select(t.id) + .where(t.segment_id == ParameterValue(_db.uuid_to_db(_id))) + ) + ) + ) + sql, params = get_sql(q_fts) + with _db.tx() as cur: + res = cur.execute(sql, params) + # assert that all FTS rows are gone + assert len(res.fetchall()) == 0 + + +def test_delete_single_fts_record( + system: System, + sample_embeddings: Iterator[OperationRecord], + produce_fns: ProducerFn, +) -> None: + producer = system.instance(Producer) + system.reset_state() + collection_id = segment_definition["collection"] + # We know that the collection_id exists so we can cast + collection_id = cast(uuid.UUID, collection_id) + + segment = SqliteMetadataSegment(system, segment_definition) + segment.start() + + embeddings, seq_ids = produce_fns(producer, collection_id, sample_embeddings, 10) + max_id = seq_ids[-1] + + sync(segment, max_id) + + assert segment.count() == 10 + results = segment.get_metadata(ids=["embedding_0"]) + assert_equiv_records(embeddings[:1], results) + _id = segment._id + _db = system.instance(SqliteDB) + # Delete by ID + delete_embedding = OperationRecord( + id="embedding_0", + embedding=None, + encoding=None, + metadata=None, + operation=Operation.DELETE, + ) + max_id = produce_fns( + producer, collection_id, (delete_embedding for _ in range(1)), 1 + )[1][-1] + t = Table("embeddings") + + sync(segment, max_id) + fts_t = Table("embedding_fulltext_search") + q_fts = ( + _db.querybuilder() + .from_(fts_t) + .select() + .where( + fts_t.rowid.isin( + _db.querybuilder() + .from_(t) + .select(t.id) + .where(t.segment_id == ParameterValue(_db.uuid_to_db(_id))) + .where(t.embedding_id == ParameterValue(delete_embedding["id"])) + ) + ) + ) + sql, params = get_sql(q_fts) + with _db.tx() as cur: + res = cur.execute(sql, params) + # assert that the ids that are deleted from the segment are also gone from the fts table + assert len(res.fetchall()) == 0 + + +def test_metadata_validation_forbidden_key() -> None: + with pytest.raises(ValueError, match="chroma:document"): + validate_metadata( + {"chroma:document": "this is not the document you are looking for"} + ) diff --git a/chromadb/test/segment/test_vector.py b/chromadb/test/segment/test_vector.py index 1ba9802c66f..e33fde65604 100644 --- a/chromadb/test/segment/test_vector.py +++ b/chromadb/test/segment/test_vector.py @@ -3,7 +3,7 @@ from chromadb.config import System, Settings from chromadb.test.conftest import ProducerFn from chromadb.types import ( - SubmitEmbeddingRecord, + OperationRecord, VectorQuery, Operation, ScalarEncoding, @@ -78,20 +78,19 @@ def system(request: FixtureRequest) -> Generator[System, None, None]: @pytest.fixture(scope="function") -def sample_embeddings() -> Iterator[SubmitEmbeddingRecord]: +def sample_embeddings() -> Iterator[OperationRecord]: """Generate a sequence of embeddings with the property that for each embedding (other than the first and last), it's nearest neighbor is the previous in the sequence, and it's second nearest neighbor is the subsequent""" - def create_record(i: int) -> SubmitEmbeddingRecord: + def create_record(i: int) -> OperationRecord: vector = [i**1.1, i**1.1] - record = SubmitEmbeddingRecord( + record = OperationRecord( id=f"embedding_{i}", embedding=vector, encoding=ScalarEncoding.FLOAT32, metadata=None, operation=Operation.ADD, - collection_id=uuid.UUID(int=0), ) return record @@ -112,8 +111,7 @@ def create_random_segment_definition() -> Segment: id=uuid.uuid4(), type="test_type", scope=SegmentScope.VECTOR, - topic="persistent://test/test/test_topic_1", - collection=None, + collection=uuid.UUID(int=0), metadata=test_hnsw_config, ) @@ -130,7 +128,7 @@ def sync(segment: VectorReader, seq_id: SeqId) -> None: def test_insert_and_count( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: @@ -138,10 +136,15 @@ def test_insert_and_count( system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) max_id = produce_fns( - producer=producer, topic=topic, n=3, embeddings=sample_embeddings + producer=producer, + collection_id=collection_id, + n=3, + embeddings=sample_embeddings, )[1][-1] segment = vector_reader(system, segment_definition) @@ -152,7 +155,10 @@ def test_insert_and_count( assert segment.count() == 3 max_id = produce_fns( - producer=producer, topic=topic, n=3, embeddings=sample_embeddings + producer=producer, + collection_id=collection_id, + n=3, + embeddings=sample_embeddings, )[1][-1] sync(segment, max_id) @@ -169,20 +175,25 @@ def approx_equal_vector(a: Vector, b: Vector, epsilon: float = 0.0001) -> bool: def test_get_vectors( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() embeddings, seq_ids = produce_fns( - producer=producer, topic=topic, embeddings=sample_embeddings, n=10 + producer=producer, + collection_id=collection_id, + embeddings=sample_embeddings, + n=10, ) sync(segment, seq_ids[-1]) @@ -196,7 +207,6 @@ def test_get_vectors( assert approx_equal_vector( actual["embedding"], cast(Vector, expected["embedding"]) ) - assert actual["seq_id"] == seq_id # Get selected IDs ids = [e["id"] for e in embeddings[5:]] @@ -208,25 +218,29 @@ def test_get_vectors( assert approx_equal_vector( actual["embedding"], cast(Vector, expected["embedding"]) ) - assert actual["seq_id"] == seq_id def test_ann_query( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() embeddings, seq_ids = produce_fns( - producer=producer, topic=topic, embeddings=sample_embeddings, n=100 + producer=producer, + collection_id=collection_id, + embeddings=sample_embeddings, + n=100, ) sync(segment, seq_ids[-1]) @@ -277,38 +291,42 @@ def test_ann_query( def test_delete( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() embeddings, seq_ids = produce_fns( - producer=producer, topic=topic, embeddings=sample_embeddings, n=5 + producer=producer, + collection_id=collection_id, + embeddings=sample_embeddings, + n=5, ) sync(segment, seq_ids[-1]) assert segment.count() == 5 - delete_record = SubmitEmbeddingRecord( + delete_record = OperationRecord( id=embeddings[0]["id"], embedding=None, encoding=None, metadata=None, operation=Operation.DELETE, - collection_id=uuid.UUID(int=0), ) assert isinstance(seq_ids, List) seq_ids.append( produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, n=1, embeddings=(delete_record for _ in range(1)), )[1][0] @@ -344,7 +362,7 @@ def test_delete( seq_ids.append( produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, n=1, embeddings=(delete_record for _ in range(1)), )[1][0] @@ -357,9 +375,9 @@ def test_delete( def _test_update( producer: Producer, - topic: str, + collection_id: uuid.UUID, segment: VectorReader, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], operation: Operation, ) -> None: """Tests the common code paths between update & upsert""" @@ -368,21 +386,20 @@ def _test_update( seq_ids: List[SeqId] = [] for e in embeddings: - seq_ids.append(producer.submit_embedding(topic, e)) + seq_ids.append(producer.submit_embedding(collection_id, e)) sync(segment, seq_ids[-1]) assert segment.count() == 3 seq_ids.append( producer.submit_embedding( - topic, - SubmitEmbeddingRecord( + collection_id, + OperationRecord( id=embeddings[0]["id"], embedding=[10.0, 10.0], encoding=ScalarEncoding.FLOAT32, metadata=None, operation=operation, - collection_id=uuid.UUID(int=0), ), ) ) @@ -419,32 +436,33 @@ def _test_update( def test_update( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() - _test_update(producer, topic, segment, sample_embeddings, Operation.UPDATE) + _test_update(producer, collection_id, segment, sample_embeddings, Operation.UPDATE) # test updating a nonexistent record - update_record = SubmitEmbeddingRecord( + update_record = OperationRecord( id="no_such_record", embedding=[10.0, 10.0], encoding=ScalarEncoding.FLOAT32, metadata=None, operation=Operation.UPDATE, - collection_id=uuid.UUID(int=0), ) seq_id = produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, n=1, embeddings=(update_record for _ in range(1)), )[1][0] @@ -457,32 +475,33 @@ def test_update( def test_upsert( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() - _test_update(producer, topic, segment, sample_embeddings, Operation.UPSERT) + _test_update(producer, collection_id, segment, sample_embeddings, Operation.UPSERT) # test updating a nonexistent record - upsert_record = SubmitEmbeddingRecord( + upsert_record = OperationRecord( id="no_such_record", embedding=[42, 42], encoding=ScalarEncoding.FLOAT32, metadata=None, operation=Operation.UPSERT, - collection_id=uuid.UUID(int=0), ) seq_id = produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, n=1, embeddings=(upsert_record for _ in range(1)), )[1][0] @@ -502,62 +521,67 @@ def test_delete_without_add( producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() assert segment.count() == 0 - delete_record = SubmitEmbeddingRecord( + delete_record = OperationRecord( id="not_in_db", embedding=None, encoding=None, metadata=None, operation=Operation.DELETE, - collection_id=uuid.UUID(int=0), ) try: - producer.submit_embedding(topic, delete_record) + producer.submit_embedding(collection_id, delete_record) except BaseException: pytest.fail("Unexpected error. Deleting on an empty segment should not raise.") def test_delete_with_local_segment_storage( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() embeddings, seq_ids = produce_fns( - producer=producer, topic=topic, embeddings=sample_embeddings, n=5 + producer=producer, + collection_id=collection_id, + embeddings=sample_embeddings, + n=5, ) sync(segment, seq_ids[-1]) assert segment.count() == 5 - delete_record = SubmitEmbeddingRecord( + delete_record = OperationRecord( id=embeddings[0]["id"], embedding=None, encoding=None, metadata=None, operation=Operation.DELETE, - collection_id=uuid.UUID(int=0), ) assert isinstance(seq_ids, List) seq_ids.append( produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, n=1, embeddings=(delete_record for _ in range(1)), )[1][0] @@ -602,38 +626,42 @@ def test_delete_with_local_segment_storage( def test_reset_state_ignored_for_allow_reset_false( system: System, - sample_embeddings: Iterator[SubmitEmbeddingRecord], + sample_embeddings: Iterator[OperationRecord], vector_reader: Type[VectorReader], produce_fns: ProducerFn, ) -> None: producer = system.instance(Producer) system.reset_state() segment_definition = create_random_segment_definition() - topic = str(segment_definition["topic"]) + collection_id = segment_definition["collection"] + # We know that the segment definition has a collection_id + collection_id = cast(uuid.UUID, collection_id) segment = vector_reader(system, segment_definition) segment.start() embeddings, seq_ids = produce_fns( - producer=producer, topic=topic, embeddings=sample_embeddings, n=5 + producer=producer, + collection_id=collection_id, + embeddings=sample_embeddings, + n=5, ) sync(segment, seq_ids[-1]) assert segment.count() == 5 - delete_record = SubmitEmbeddingRecord( + delete_record = OperationRecord( id=embeddings[0]["id"], embedding=None, encoding=None, metadata=None, operation=Operation.DELETE, - collection_id=uuid.UUID(int=0), ) assert isinstance(seq_ids, List) seq_ids.append( produce_fns( producer=producer, - topic=topic, + collection_id=collection_id, n=1, embeddings=(delete_record for _ in range(1)), )[1][0] diff --git a/chromadb/test/test_api.py b/chromadb/test/test_api.py index 36a82205e45..cb88ed2bb77 100644 --- a/chromadb/test/test_api.py +++ b/chromadb/test/test_api.py @@ -1,5 +1,7 @@ # type: ignore +import traceback import requests +from urllib3.connectionpool import InsecureRequestWarning import chromadb from chromadb.api.fastapi import FastAPI @@ -360,6 +362,7 @@ def test_modify_error_on_existing_name(api): with pytest.raises(Exception): c2.modify(name="testspace") + def test_modify_warn_on_DF_change(api, caplog): api.reset() @@ -368,6 +371,7 @@ def test_modify_warn_on_DF_change(api, caplog): with pytest.raises(Exception, match="not supported") as e: collection.modify(metadata={"hnsw:space": "cosine"}) + def test_metadata_cru(api): api.reset() metadata_a = {"a": 1, "b": 2} @@ -1437,6 +1441,7 @@ def test_invalid_embeddings(api): # test to make sure update shows exception for bad dimensionality + def test_dimensionality_exception_update(api): api.reset() collection = api.create_collection("test_dimensionality_update_exception") @@ -1446,8 +1451,10 @@ def test_dimensionality_exception_update(api): collection.update(**bad_dimensionality_records) assert "dimensionality" in str(e.value) + # test to make sure upsert shows exception for bad dimensionality + def test_dimensionality_exception_upsert(api): api.reset() collection = api.create_collection("test_dimensionality_upsert_exception") @@ -1456,3 +1463,39 @@ def test_dimensionality_exception_upsert(api): with pytest.raises(Exception) as e: collection.upsert(**bad_dimensionality_records) assert "dimensionality" in str(e.value) + + +def test_ssl_self_signed(client_ssl): + if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"): + pytest.skip("Skipping test for integration test") + client_ssl.heartbeat() + + +def test_ssl_self_signed_without_ssl_verify(client_ssl): + if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"): + pytest.skip("Skipping test for integration test") + client_ssl.heartbeat() + _port = client_ssl._server._settings.chroma_server_http_port + with pytest.raises(ValueError) as e: + chromadb.HttpClient(ssl=True, port=_port) + stack_trace = traceback.format_exception( + type(e.value), e.value, e.value.__traceback__ + ) + client_ssl.clear_system_cache() + assert "CERTIFICATE_VERIFY_FAILED" in "".join(stack_trace) + + +def test_ssl_self_signed_with_verify_false(client_ssl): + if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"): + pytest.skip("Skipping test for integration test") + client_ssl.heartbeat() + _port = client_ssl._server._settings.chroma_server_http_port + with pytest.warns(InsecureRequestWarning) as record: + client = chromadb.HttpClient( + ssl=True, + port=_port, + settings=chromadb.Settings(chroma_server_ssl_verify=False), + ) + client.heartbeat() + client_ssl.clear_system_cache() + assert "Unverified HTTPS request" in str(record[0].message) diff --git a/chromadb/test/test_chroma.py b/chromadb/test/test_chroma.py index 9d88ea8cc49..89b4ae924eb 100644 --- a/chromadb/test/test_chroma.py +++ b/chromadb/test/test_chroma.py @@ -66,7 +66,7 @@ def test_fastapi(self, mock: Mock) -> None: chroma_api_impl="chromadb.api.fastapi.FastAPI", persist_directory="./foo", chroma_server_host="foo", - chroma_server_http_port="80", + chroma_server_http_port=80, ) ) assert mock.called @@ -78,7 +78,7 @@ def test_settings_pass_to_fastapi(self, mock: Mock) -> None: settings = chromadb.config.Settings( chroma_api_impl="chromadb.api.fastapi.FastAPI", chroma_server_host="foo", - chroma_server_http_port="80", + chroma_server_http_port=80, chroma_server_headers={"foo": "bar"}, ) client = chromadb.Client(settings) @@ -106,7 +106,7 @@ def test_legacy_values() -> None: chroma_api_impl="chromadb.api.local.LocalAPI", persist_directory="./foo", chroma_server_host="foo", - chroma_server_http_port="80", + chroma_server_http_port=80, ) ) client.clear_system_cache() diff --git a/chromadb/test/test_client.py b/chromadb/test/test_client.py index f67293d8586..34dd2df1412 100644 --- a/chromadb/test/test_client.py +++ b/chromadb/test/test_client.py @@ -60,9 +60,9 @@ def test_http_client_with_inconsistent_host_settings() -> None: def test_http_client_with_inconsistent_port_settings() -> None: try: chromadb.HttpClient( - port="8002", + port=8002, settings=Settings( - chroma_server_http_port="8001", + chroma_server_http_port=8001, ), ) except ValueError as e: diff --git a/chromadb/test/test_logservice.py b/chromadb/test/test_logservice.py new file mode 100644 index 00000000000..8cac506c16c --- /dev/null +++ b/chromadb/test/test_logservice.py @@ -0,0 +1,156 @@ +import array +from typing import Dict, Any, Callable + +from chromadb.config import System, Settings +from chromadb.logservice.logservice import LogService +from chromadb.test.conftest import skip_if_not_cluster +from chromadb.test.test_api import records # type: ignore +from chromadb.api.models.Collection import Collection + +batch_records = { + "embeddings": [[1.1, 2.3, 3.2], [1.2, 2.24, 3.2]], + "ids": ["https://example.com/1", "https://example.com/2"], +} + +metadata_records = { + "embeddings": [[1.1, 2.3, 3.2], [1.2, 2.24, 3.2]], + "ids": ["id1", "id2"], + "metadatas": [ + {"int_value": 1, "string_value": "one", "float_value": 1.001}, + {"int_value": 2}, + ], +} + +contains_records = { + "embeddings": [[1.1, 2.3, 3.2], [1.2, 2.24, 3.2]], + "documents": ["this is doc1 and it's great!", "doc2 is also great!"], + "ids": ["id1", "id2"], + "metadatas": [ + {"int_value": 1, "string_value": "one", "float_value": 1.001}, + {"int_value": 2, "float_value": 2.002, "string_value": "two"}, + ], +} + + +def verify_records( + logservice: LogService, + collection: Collection, + test_records_map: Dict[str, Dict[str, Any]], + test_func: Callable, # type: ignore + operation: int, +) -> None: + start_offset = 1 + for batch_records in test_records_map.values(): + test_func(**batch_records) + pushed_records = logservice.pull_logs(collection.id, start_offset, 100) + assert len(pushed_records) == len(batch_records["ids"]) + for i, record in enumerate(pushed_records): + assert record.record.id == batch_records["ids"][i] + assert record.record.operation == operation + embedding = array.array("f", batch_records["embeddings"][i]).tobytes() + assert record.record.vector.vector == embedding + metadata_count = 0 + if "metadatas" in batch_records: + metadata_count += len(batch_records["metadatas"][i]) + for key, value in batch_records["metadatas"][i].items(): + if isinstance(value, int): + assert record.record.metadata.metadata[key].int_value == value + elif isinstance(value, float): + assert record.record.metadata.metadata[key].float_value == value + elif isinstance(value, str): + assert ( + record.record.metadata.metadata[key].string_value == value + ) + else: + assert False + if "documents" in batch_records: + metadata_count += 1 + assert ( + record.record.metadata.metadata["chroma:document"].string_value + == batch_records["documents"][i] + ) + assert len(record.record.metadata.metadata) == metadata_count + start_offset += len(pushed_records) + + +@skip_if_not_cluster() +def test_add(api): # type: ignore + system = System(Settings(allow_reset=True)) + logservice = system.instance(LogService) + system.start() + api.reset() + + test_records_map = { + "batch_records": batch_records, + "metadata_records": metadata_records, + "contains_records": contains_records, + } + + collection = api.create_collection("testadd") + verify_records(logservice, collection, test_records_map, collection.add, 0) + + +@skip_if_not_cluster() +def test_update(api): # type: ignore + system = System(Settings(allow_reset=True)) + logservice = system.instance(LogService) + system.start() + api.reset() + + test_records_map = { + "updated_records": { + "ids": [records["ids"][0]], + "embeddings": [[0.1, 0.2, 0.3]], + "metadatas": [{"foo": "bar"}], + }, + } + + collection = api.create_collection("testupdate") + verify_records(logservice, collection, test_records_map, collection.update, 1) + + +@skip_if_not_cluster() +def test_delete(api): # type: ignore + system = System(Settings(allow_reset=True)) + logservice = system.instance(LogService) + system.start() + api.reset() + + collection = api.create_collection("testdelete") + + # push 2 records + collection.add(**contains_records) + pushed_records = logservice.pull_logs(collection.id, 1, 100) + assert len(pushed_records) == 2 + + # delete by where does not work atm + collection.delete(where_document={"$contains": "doc1"}) + collection.delete(where_document={"$contains": "bad"}) + collection.delete(where_document={"$contains": "great"}) + pushed_records = logservice.pull_logs(collection.id, 3, 100) + assert len(pushed_records) == 0 + + # delete by ids + collection.delete(ids=["id1", "id2"]) + pushed_records = logservice.pull_logs(collection.id, 3, 100) + assert len(pushed_records) == 2 + for record in pushed_records: + assert record.record.operation == 3 + assert record.record.id in ["id1", "id2"] + + +@skip_if_not_cluster() +def test_upsert(api): # type: ignore + system = System(Settings(allow_reset=True)) + logservice = system.instance(LogService) + system.start() + api.reset() + + test_records_map = { + "batch_records": batch_records, + "metadata_records": metadata_records, + "contains_records": contains_records, + } + + collection = api.create_collection("testupsert") + verify_records(logservice, collection, test_records_map, collection.upsert, 2) diff --git a/chromadb/test/utils/test_messagid.py b/chromadb/test/utils/test_messagid.py index eff20a1b6fe..64d80e9b6b0 100644 --- a/chromadb/test/utils/test_messagid.py +++ b/chromadb/test/utils/test_messagid.py @@ -1,93 +1,19 @@ import chromadb.utils.messageid as mid -import pulsar import hypothesis.strategies as st -from hypothesis import given, settings, note -from typing import Any, Tuple +from hypothesis import given, settings @st.composite -def message_id(draw: st.DrawFn) -> pulsar.MessageId: - ledger_id = draw(st.integers(min_value=0, max_value=2**63 - 1)) - entry_id = draw(st.integers(min_value=0, max_value=2**63 - 1)) - batch_index = draw(st.integers(min_value=(2**31 - 1) * -1, max_value=2**31 - 1)) - partition = draw(st.integers(min_value=(2**31 - 1) * -1, max_value=2**31 - 1)) - return pulsar.MessageId(partition, ledger_id, entry_id, batch_index) +def message_id(draw: st.DrawFn) -> int: + offset_id = draw(st.integers(min_value=0, max_value=2**63 - 1)) + return offset_id @given(message_id=message_id()) @settings(max_examples=10000) # these are very fast and we want good coverage -def test_roundtrip_formats(message_id: pulsar.MessageId) -> None: - int1 = mid.pulsar_to_int(message_id) - - # Roundtrip int->string and back - str1 = mid.int_to_str(int1) - assert int1 == mid.str_to_int(str1) +def test_roundtrip_formats(message_id: int) -> None: + int1 = message_id # Roundtrip int->bytes and back b1 = mid.int_to_bytes(int1) assert int1 == mid.bytes_to_int(b1) - - # Roundtrip int -> MessageId and back - message_id_result = mid.int_to_pulsar(int1) - assert message_id_result.partition() == message_id.partition() - assert message_id_result.ledger_id() == message_id.ledger_id() - assert message_id_result.entry_id() == message_id.entry_id() - assert message_id_result.batch_index() == message_id.batch_index() - - -def assert_compare(pair1: Tuple[Any, Any], pair2: Tuple[Any, Any]) -> None: - """Helper function: assert that the two pairs of values always compare in the same - way across all comparisons and orderings.""" - - a, b = pair1 - c, d = pair2 - - try: - assert (a > b) == (c > d) - assert (a >= b) == (c >= d) - assert (a < b) == (c < d) - assert (a <= b) == (c <= d) - assert (a == b) == (c == d) - except AssertionError: - note(f"Failed to compare {a} and {b} with {c} and {d}") - note(f"type: {type(a)}") - raise - - -@given(m1=message_id(), m2=message_id()) -@settings(max_examples=10000) # these are very fast and we want good coverage -def test_messageid_comparison(m1: pulsar.MessageId, m2: pulsar.MessageId) -> None: - # MessageID comparison is broken in the Pulsar Python & CPP libraries: - # The partition field is not taken into account, and two MessageIDs with different - # partitions will compare inconsistently (m1 > m2 AND m2 > m1) - # To avoid this, we zero-out the partition field before testing. - m1 = pulsar.MessageId(0, m1.ledger_id(), m1.entry_id(), m1.batch_index()) - m2 = pulsar.MessageId(0, m2.ledger_id(), m2.entry_id(), m2.batch_index()) - - i1 = mid.pulsar_to_int(m1) - i2 = mid.pulsar_to_int(m2) - - # In python, MessageId objects are not comparable directory, but the - # internal generated native object is. - internal1 = m1._msg_id - internal2 = m2._msg_id - - s1 = mid.int_to_str(i1) - s2 = mid.int_to_str(i2) - - # assert that all strings, all ints, and all native objects compare the same - assert_compare((internal1, internal2), (i1, i2)) - assert_compare((internal1, internal2), (s1, s2)) - - -def test_max_values() -> None: - pulsar.MessageId(2**31 - 1, 2**63 - 1, 2**63 - 1, 2**31 - 1) - - -@given( - i1=st.integers(min_value=0, max_value=2**192 - 1), - i2=st.integers(min_value=0, max_value=2**192 - 1), -) -@settings(max_examples=10000) # these are very fast and we want good coverage -def test_string_comparison(i1: int, i2: int) -> None: - assert_compare((i1, i2), (mid.int_to_str(i1), mid.int_to_str(i2))) diff --git a/chromadb/types.py b/chromadb/types.py index fd66f12af6c..205a01fbb55 100644 --- a/chromadb/types.py +++ b/chromadb/types.py @@ -26,7 +26,6 @@ class SegmentScope(Enum): class Collection(TypedDict): id: UUID name: str - topic: str metadata: Optional[Metadata] dimension: Optional[int] tenant: str @@ -47,9 +46,6 @@ class Segment(TypedDict): id: UUID type: NamespacedName scope: SegmentScope - # If a segment has a topic, it implies that this segment is a consumer of the topic - # and indexes the contents of the topic. - topic: Optional[str] # If a segment has a collection, it implies that this segment implements the full # collection and can be used to service queries (for it's given scope.) collection: Optional[UUID] @@ -57,9 +53,9 @@ class Segment(TypedDict): # SeqID can be one of three types of value in our current and future plans: -# 1. A Pulsar MessageID encoded as a 192-bit integer -# 2. A Pulsar MessageIndex (a 64-bit integer) -# 3. A SQL RowID (a 64-bit integer) +# 1. A Pulsar MessageID encoded as a 192-bit integer - This is no longer used as we removed pulsar +# 2. A Pulsar MessageIndex (a 64-bit integer) - This is no longer used as we removed pulsar +# 3. A SQL RowID (a 64-bit integer) - This is used by both sqlite and the new log-service # All three of these types can be expressed as a Python int, so that is the type we # use in the internal Python API. However, care should be taken that the larger 192-bit @@ -79,42 +75,25 @@ class Operation(Enum): class VectorEmbeddingRecord(TypedDict): id: str - seq_id: SeqId embedding: Vector class MetadataEmbeddingRecord(TypedDict): id: str - seq_id: SeqId metadata: Optional[Metadata] -class EmbeddingRecord(TypedDict): +class OperationRecord(TypedDict): id: str - seq_id: SeqId embedding: Optional[Vector] encoding: Optional[ScalarEncoding] metadata: Optional[UpdateMetadata] operation: Operation - # The collection the operation is being performed on - # This is optional because in the single node version, - # topics are 1:1 with collections. So consumers of the ingest queue - # implicitly know this mapping. However, in the multi-node version, - # topics are shared between collections, so we need to explicitly - # specify the collection. - # For backwards compatability reasons, we can't make this a required field on - # single node, since data written with older versions of the code won't be able to - # populate it. - collection_id: Optional[UUID] - - -class SubmitEmbeddingRecord(TypedDict): - id: str - embedding: Optional[Vector] - encoding: Optional[ScalarEncoding] - metadata: Optional[UpdateMetadata] - operation: Operation - collection_id: UUID # The collection the operation is being performed on + + +class LogRecord(TypedDict): + log_offset: int + record: OperationRecord class VectorQuery(TypedDict): @@ -131,7 +110,6 @@ class VectorQueryResult(TypedDict): """A KNN/ANN query result""" id: str - seq_id: SeqId distance: float embedding: Optional[Vector] diff --git a/chromadb/utils/data_loaders.py b/chromadb/utils/data_loaders.py index 60057e0e584..82ea894aa9a 100644 --- a/chromadb/utils/data_loaders.py +++ b/chromadb/utils/data_loaders.py @@ -1,8 +1,8 @@ import importlib import multiprocessing -from typing import Optional, Sequence, List +from typing import Optional, Sequence, List, Tuple import numpy as np -from chromadb.api.types import URI, DataLoader, Image +from chromadb.api.types import URI, DataLoader, Image, URIs from concurrent.futures import ThreadPoolExecutor @@ -22,3 +22,10 @@ def _load_image(self, uri: Optional[URI]) -> Optional[Image]: def __call__(self, uris: Sequence[Optional[URI]]) -> List[Optional[Image]]: with ThreadPoolExecutor(max_workers=self._max_workers) as executor: return list(executor.map(self._load_image, uris)) + + +class ChromaLangchainPassthroughDataLoader(DataLoader[List[Optional[Image]]]): + # This is a simple pass through data loader that just returns the input data with "images" + # flag which lets the langchain embedding function know that the data is image uris + def __call__(self, uris: URIs) -> Tuple[str, URIs]: # type: ignore + return ("images", uris) diff --git a/chromadb/utils/directory.py b/chromadb/utils/directory.py new file mode 100644 index 00000000000..d470a810ed5 --- /dev/null +++ b/chromadb/utils/directory.py @@ -0,0 +1,21 @@ +import os + +def get_directory_size(directory: str) -> int: + """ + Calculate the total size of the directory by walking through each file. + + Parameters: + directory (str): The path of the directory for which to calculate the size. + + Returns: + total_size (int): The total size of the directory in bytes. + """ + total_size = 0 + for dirpath, _, filenames in os.walk(directory): + for f in filenames: + fp = os.path.join(dirpath, f) + # skip if it is symbolic link + if not os.path.islink(fp): + total_size += os.path.getsize(fp) + + return total_size \ No newline at end of file diff --git a/chromadb/utils/embedding_functions.py b/chromadb/utils/embedding_functions.py index ec5fc05e3ee..3f0a1ce043b 100644 --- a/chromadb/utils/embedding_functions.py +++ b/chromadb/utils/embedding_functions.py @@ -16,6 +16,7 @@ is_document, ) +from io import BytesIO from pathlib import Path import os import tarfile @@ -27,6 +28,7 @@ import inspect import json import sys +import base64 try: from chromadb.is_thin_client import is_thin_client @@ -61,7 +63,16 @@ def __init__( model_name: str = "all-MiniLM-L6-v2", device: str = "cpu", normalize_embeddings: bool = False, + **kwargs: Any, ): + """Initialize SentenceTransformerEmbeddingFunction. + + Args: + model_name (str, optional): Identifier of the SentenceTransformer model, defaults to "all-MiniLM-L6-v2" + device (str, optional): Device used for computation, defaults to "cpu" + normalize_embeddings (bool, optional): Whether to normalize returned vectors, defaults to False + **kwargs: Additional arguments to pass to the SentenceTransformer model. + """ if model_name not in self.models: try: from sentence_transformers import SentenceTransformer @@ -69,7 +80,9 @@ def __init__( raise ValueError( "The sentence_transformers python package is not installed. Please install it with `pip install sentence_transformers`" ) - self.models[model_name] = SentenceTransformer(model_name, device=device) + self.models[model_name] = SentenceTransformer( + model_name, device=device, **kwargs + ) self._model = self.models[model_name] self._normalize_embeddings = normalize_embeddings @@ -506,11 +519,17 @@ def model(self) -> "InferenceSession": raise ValueError( f"Preferred providers must be subset of available providers: {self.ort.get_available_providers()}" ) + + # Suppress onnxruntime warnings. This produces logspew, mainly when onnx tries to use CoreML, which doesn't fit this model. + so = self.ort.SessionOptions() + so.log_severity_level = 3 + return self.ort.InferenceSession( os.path.join(self.DOWNLOAD_PATH, self.EXTRACTED_FOLDER_NAME, "model.onnx"), # Since 1.9 onnyx runtime requires providers to be specified when there are multiple available - https://onnxruntime.ai/docs/api/python/api_summary.html # This is probably not ideal but will improve DX as no exceptions will be raised in multi-provider envs providers=self._preferred_providers, + sess_options=so, ) def __call__(self, input: Documents) -> Embeddings: @@ -665,7 +684,10 @@ def __call__(self, input: Documents) -> Embeddings: class OpenCLIPEmbeddingFunction(EmbeddingFunction[Union[Documents, Images]]): def __init__( - self, model_name: str = "ViT-B-32", checkpoint: str = "laion2b_s34b_b79k" + self, + model_name: str = "ViT-B-32", + checkpoint: str = "laion2b_s34b_b79k", + device: Optional[str] = "cpu", ) -> None: try: import open_clip @@ -691,6 +713,7 @@ def __init__( model_name=model_name, pretrained=checkpoint ) self._model = model + self._model.to(device) self._preprocess = preprocess self._tokenizer = open_clip.get_tokenizer(model_name=model_name) @@ -719,6 +742,74 @@ def __call__(self, input: Union[Documents, Images]) -> Embeddings: return embeddings +class RoboflowEmbeddingFunction(EmbeddingFunction[Union[Documents, Images]]): + def __init__( + self, api_key: str = "", api_url = "https://infer.roboflow.com" + ) -> None: + """ + Create a RoboflowEmbeddingFunction. + + Args: + api_key (str): Your API key for the Roboflow API. + api_url (str, optional): The URL of the Roboflow API. Defaults to "https://infer.roboflow.com". + """ + if not api_key: + api_key = os.environ.get("ROBOFLOW_API_KEY") + + self._api_url = api_url + self._api_key = api_key + + try: + self._PILImage = importlib.import_module("PIL.Image") + except ImportError: + raise ValueError( + "The PIL python package is not installed. Please install it with `pip install pillow`" + ) + + def __call__(self, input: Union[Documents, Images]) -> Embeddings: + embeddings = [] + + for item in input: + if is_image(item): + image = self._PILImage.fromarray(item) + + buffer = BytesIO() + image.save(buffer, format="JPEG") + base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8") + + infer_clip_payload = { + "image": { + "type": "base64", + "value": base64_image, + }, + } + + res = requests.post( + f"{self._api_url}/clip/embed_image?api_key={self._api_key}", + json=infer_clip_payload, + ) + + result = res.json()['embeddings'] + + embeddings.append(result[0]) + + elif is_document(item): + infer_clip_payload = { + "text": input, + } + + res = requests.post( + f"{self._api_url}/clip/embed_text?api_key={self._api_key}", + json=infer_clip_payload, + ) + + result = res.json()['embeddings'] + + embeddings.append(result[0]) + + return embeddings + + class AmazonBedrockEmbeddingFunction(EmbeddingFunction[Documents]): def __init__( self, @@ -809,6 +900,125 @@ def __call__(self, input: Documents) -> Embeddings: ) +def create_langchain_embedding(langchain_embdding_fn: Any): # type: ignore + try: + from langchain_core.embeddings import Embeddings as LangchainEmbeddings + except ImportError: + raise ValueError( + "The langchain_core python package is not installed. Please install it with `pip install langchain-core`" + ) + + class ChromaLangchainEmbeddingFunction( + LangchainEmbeddings, EmbeddingFunction[Union[Documents, Images]] # type: ignore + ): + """ + This class is used as bridge between langchain embedding functions and custom chroma embedding functions. + """ + + def __init__(self, embedding_function: LangchainEmbeddings) -> None: + """ + Initialize the ChromaLangchainEmbeddingFunction + + Args: + embedding_function : The embedding function implementing Embeddings from langchain_core. + """ + self.embedding_function = embedding_function + + def embed_documents(self, documents: Documents) -> List[List[float]]: + return self.embedding_function.embed_documents(documents) # type: ignore + + def embed_query(self, query: str) -> List[float]: + return self.embedding_function.embed_query(query) # type: ignore + + def embed_image(self, uris: List[str]) -> List[List[float]]: + if hasattr(self.embedding_function, "embed_image"): + return self.embedding_function.embed_image(uris) # type: ignore + else: + raise ValueError( + "The provided embedding function does not support image embeddings." + ) + + def __call__(self, input: Documents) -> Embeddings: # type: ignore + """ + Get the embeddings for a list of texts or images. + + Args: + input (Documents | Images): A list of texts or images to get embeddings for. + Images should be provided as a list of URIs passed through the langchain data loader + + Returns: + Embeddings: The embeddings for the texts or images. + + Example: + >>> langchain_embedding = ChromaLangchainEmbeddingFunction(embedding_function=OpenAIEmbeddings(model="text-embedding-3-large")) + >>> texts = ["Hello, world!", "How are you?"] + >>> embeddings = langchain_embedding(texts) + """ + # Due to langchain quirks, the dataloader returns a tuple if the input is uris of images + if input[0] == "images": + return self.embed_image(list(input[1])) # type: ignore + + return self.embed_documents(list(input)) # type: ignore + + return ChromaLangchainEmbeddingFunction(embedding_function=langchain_embdding_fn) + + +class OllamaEmbeddingFunction(EmbeddingFunction[Documents]): + """ + This class is used to generate embeddings for a list of texts using the Ollama Embedding API (https://github.com/ollama/ollama/blob/main/docs/api.md#generate-embeddings). + """ + + def __init__(self, url: str, model_name: str) -> None: + """ + Initialize the Ollama Embedding Function. + + Args: + url (str): The URL of the Ollama Server. + model_name (str): The name of the model to use for text embeddings. E.g. "nomic-embed-text" (see https://ollama.com/library for available models). + """ + try: + import requests + except ImportError: + raise ValueError( + "The requests python package is not installed. Please install it with `pip install requests`" + ) + self._api_url = f"{url}" + self._model_name = model_name + self._session = requests.Session() + + def __call__(self, input: Documents) -> Embeddings: + """ + Get the embeddings for a list of texts. + + Args: + input (Documents): A list of texts to get embeddings for. + + Returns: + Embeddings: The embeddings for the texts. + + Example: + >>> ollama_ef = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings", model_name="nomic-embed-text") + >>> texts = ["Hello, world!", "How are you?"] + >>> embeddings = ollama_ef(texts) + """ + # Call Ollama Server API for each document + texts = input if isinstance(input, list) else [input] + embeddings = [ + self._session.post( + self._api_url, json={"model": self._model_name, "prompt": text} + ).json() + for text in texts + ] + return cast( + Embeddings, + [ + embedding["embedding"] + for embedding in embeddings + if "embedding" in embedding + ], + ) + + # List of all classes in this module _classes = [ name diff --git a/chromadb/server/fastapi/utils.py b/chromadb/utils/fastapi.py similarity index 98% rename from chromadb/server/fastapi/utils.py rename to chromadb/utils/fastapi.py index b7e781dae68..8300880e402 100644 --- a/chromadb/server/fastapi/utils.py +++ b/chromadb/utils/fastapi.py @@ -10,8 +10,9 @@ def fastapi_json_response(error: ChromaError) -> JSONResponse: status_code=error.code(), ) + def string_to_uuid(uuid_str: str) -> UUID: try: return UUID(uuid_str) except ValueError: - raise InvalidUUIDError(f"Could not parse {uuid_str} as a UUID") \ No newline at end of file + raise InvalidUUIDError(f"Could not parse {uuid_str} as a UUID") diff --git a/chromadb/utils/messageid.py b/chromadb/utils/messageid.py index 9501f36c759..2583a7b420c 100644 --- a/chromadb/utils/messageid.py +++ b/chromadb/utils/messageid.py @@ -1,36 +1,3 @@ -import pulsar - - -def pulsar_to_int(message_id: pulsar.MessageId) -> int: - ledger_id: int = message_id.ledger_id() - entry_id: int = message_id.entry_id() - batch_index: int = message_id.batch_index() - partition: int = message_id.partition() - - # Convert to offset binary encoding to preserve ordering semantics when encoded - # see https://en.wikipedia.org/wiki/Offset_binary - ledger_id = ledger_id + 2**63 - entry_id = entry_id + 2**63 - batch_index = batch_index + 2**31 - partition = partition + 2**31 - - return ledger_id << 128 | entry_id << 64 | batch_index << 32 | partition - - -def int_to_pulsar(message_id: int) -> pulsar.MessageId: - partition = message_id & 0xFFFFFFFF - batch_index = message_id >> 32 & 0xFFFFFFFF - entry_id = message_id >> 64 & 0xFFFFFFFFFFFFFFFF - ledger_id = message_id >> 128 & 0xFFFFFFFFFFFFFFFF - - partition = partition - 2**31 - batch_index = batch_index - 2**31 - entry_id = entry_id - 2**63 - ledger_id = ledger_id - 2**63 - - return pulsar.MessageId(partition, ledger_id, entry_id, batch_index) - - def int_to_bytes(int: int) -> bytes: """Convert int to a 24 byte big endian byte string""" return int.to_bytes(24, "big") @@ -39,42 +6,3 @@ def int_to_bytes(int: int) -> bytes: def bytes_to_int(bytes: bytes) -> int: """Convert a 24 byte big endian byte string to an int""" return int.from_bytes(bytes, "big") - - -# Sorted in lexographic order -base85 = ( - "!#$%&()*+-0123456789;<=>?@ABCDEFGHIJKLMNOP" - + "QRSTUVWXYZ^_`abcdefghijklmnopqrstuvwxyz{|}~" -) - - -# not the most efficient way to do this, see benchmark function below -def _int_to_str(n: int) -> str: - if n < 85: - return base85[n] - else: - return _int_to_str(n // 85) + base85[n % 85] - - -def int_to_str(n: int) -> str: - return _int_to_str(n).rjust(36, "!") # left pad with '!' to 36 chars - - -def str_to_int(s: str) -> int: - return sum(base85.index(c) * 85**i for i, c in enumerate(s[::-1])) - - -# 1m in 5 seconds on a M1 Pro -# Not fast, but not likely to be a bottleneck either -def _benchmark() -> None: - import random - import time - - t0 = time.time() - for i in range(1000000): - x = random.randint(0, 2**192 - 1) - s = int_to_str(x) - if s == "!": # prevent compiler from optimizing out - print("oops") - t1 = time.time() - print(t1 - t0) diff --git a/clients/js/DEVELOP.md b/clients/js/DEVELOP.md index c82fd15327e..1604251477e 100644 --- a/clients/js/DEVELOP.md +++ b/clients/js/DEVELOP.md @@ -6,20 +6,34 @@ This readme is helpful for local dev. - Make sure you have Java installed (for the generator). You can download it from [java.com](https://java.com) - Make sure you set ALLOW_RESET=True for your Docker Container. If you don't do this, tests won't pass. + ``` environment: - IS_PERSISTENT=TRUE - ALLOW_RESET=True ``` + - Make sure you are running the docker backend at localhost:8000 (\*there is probably a way to stand up the fastapi server by itself and programmatically in the loop of generating this, but not prioritizing it for now. It may be important for the release) -### Generating +### Running the Examples + +To get started developing on the JS client libraries, you'll want to run the examples. + +1. `yarn` to install deps. +1. `yarn build` to build the library. +1. `cd examples/browser` or `cd examples/node` +1. `yarn` to install example deps. +1. `yarn dev` to run the example. + +### Generating REST Client Code + +If you modify the REST API, you'll need to regenerate the generated code that underlies the JavaScript client libraries. 1. `yarn` to install deps 2. `yarn genapi` 3. Examples are in the `examples` folder. There is one for the browser and one for node. Run them with `yarn dev`, eg `cd examples/browser && yarn dev` -### Running test +### Running tests `yarn test` will launch a test docker backend, run a db cleanup and run tests. `yarn test:run` will run against the docker backend you have running. But CAUTION, it will delete data. This is the easiest and fastest way to run tests. @@ -29,8 +43,9 @@ environment: #### Automatically ##### Increase the version number + 1. Create a new PR for the release that upgrades the version in code. Name it `js_release/A.B.C` for production releases and `js_release_alpha/A.B.C` for alpha releases. In the package.json update the version number to the new version. For production releases this is just the version number, for alpha -releases this is the version number with '-alphaX' appended to it. For example, if the current version is 1.0.0, the alpha release would be 1.0.0-alpha1 for the first alpha release, 1.0.0-alpha2 for the second alpha release, etc. + releases this is the version number with '-alphaX' appended to it. For example, if the current version is 1.0.0, the alpha release would be 1.0.0-alpha1 for the first alpha release, 1.0.0-alpha2 for the second alpha release, etc. 2. Add the "release" label to this PR 3. Once the PR is merged, tag your commit SHA with the release version @@ -45,6 +60,7 @@ git tag js_release_alpha_A.B.C 4. You need to then wait for the github action for main for `chroma js release` to complete on main. ##### Perform the release + 1. Push your tag to origin to create the release ```bash @@ -55,12 +71,12 @@ git push origin js_release_A.B.C git push origin js_release_alpha_A.B.C ``` + 2. This will trigger a Github action which performs the release #### Manually -`npm run release` pushes the `package.json` defined packaged to the package manager for authenticated users. It will build, test, and then publish the new version. - +`npm run release` pushes the `package.json` defined packaged to the package manager for authenticated users. It will build, test, and then publish the new version. ### Useful links diff --git a/clients/js/examples/browser/app.ts b/clients/js/examples/browser/app.ts index afc9ddbb766..47e4c02c5cf 100644 --- a/clients/js/examples/browser/app.ts +++ b/clients/js/examples/browser/app.ts @@ -1,4 +1,4 @@ -import { ChromaClient } from '../../src/ChromaClient'; +import { ChromaClient } from "../../src/ChromaClient"; // import env.ts window.onload = async () => { @@ -27,7 +27,7 @@ window.onload = async () => { const queryData = await collection.query({ queryEmbeddings: [1, 2, 3, 4, 5], nResults: 5, - where: { test: "test" } + where: { test: "test" }, }); console.log("queryData", queryData); diff --git a/clients/js/examples/node/app.js b/clients/js/examples/node/app.js index 0b153aaae35..b4ad303ab58 100644 --- a/clients/js/examples/node/app.js +++ b/clients/js/examples/node/app.js @@ -9,7 +9,9 @@ app.get("/", async (req, res) => { const cc = new chroma.ChromaClient({ path: "http://localhost:8000" }); await cc.reset(); - const google = new chroma.GoogleGenerativeAiEmbeddingFunction({ googleApiKey:"" }); + const google = new chroma.GoogleGenerativeAiEmbeddingFunction({ + googleApiKey: "", + }); const collection = await cc.createCollection({ name: "test-from-js", @@ -18,16 +20,16 @@ app.get("/", async (req, res) => { await collection.add({ ids: ["doc1", "doc2"], - documents: [ - "doc1", - "doc2", - ] + documents: ["doc1", "doc2"], }); let count = await collection.count(); console.log("count", count); - const googleQuery = new chroma.GoogleGenerativeAiEmbeddingFunction({ googleApiKey:"", taskType: 'RETRIEVAL_QUERY' }); + const googleQuery = new chroma.GoogleGenerativeAiEmbeddingFunction({ + googleApiKey: "", + taskType: "RETRIEVAL_QUERY", + }); const queryCollection = await cc.getCollection({ name: "test-from-js", @@ -36,16 +38,16 @@ app.get("/", async (req, res) => { const query = await collection.query({ queryTexts: ["doc1"], - nResults: 1 + nResults: 1, }); console.log("query", query); console.log("COMPLETED"); const collections = await cc.listCollections(); - console.log('collections', collections) + console.log("collections", collections); - res.send('Hello World!'); + res.send("Hello World!"); }); app.listen(3000, function () { console.log("Example app listening on port 3000!"); diff --git a/clients/js/package.json b/clients/js/package.json index 5fa81664bad..08ccf5da229 100644 --- a/clients/js/package.json +++ b/clients/js/package.json @@ -63,6 +63,7 @@ "db:run-auth-xtoken": "cd ../.. && echo \"CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER=X_CHROMA_TOKEN\nCHROMA_SERVER_AUTH_CREDENTIALS=test-token\nCHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=chromadb.auth.token.TokenConfigServerAuthCredentialsProvider\nCHROMA_SERVER_AUTH_PROVIDER=chromadb.auth.token.TokenAuthServerProvider\\nCHROMA_PORT=8001\" > .chroma_env && docker-compose -f docker-compose.test-auth.yml --env-file ./.chroma_env up --detach && sleep 5", "prebuild": "rimraf dist", "build": "tsup", + "watch": "tsup --watch", "genapi": "./genapi.sh", "prettier": "prettier --write .", "release": "run-s build test:run && npm publish", @@ -72,8 +73,8 @@ "node": ">=14.17.0" }, "dependencies": { - "isomorphic-fetch": "^3.0.0", - "cliui": "^8.0.1" + "cliui": "^8.0.1", + "isomorphic-fetch": "^3.0.0" }, "peerDependencies": { "@google/generative-ai": "^0.1.1", diff --git a/clients/js/src/AdminClient.ts b/clients/js/src/AdminClient.ts index 7de713e8d4e..246f181710e 100644 --- a/clients/js/src/AdminClient.ts +++ b/clients/js/src/AdminClient.ts @@ -1,272 +1,266 @@ import { Configuration, ApiApi as DefaultApi } from "./generated"; import { handleSuccess, handleError, validateTenantDatabase } from "./utils"; -import { ConfigOptions } from './types'; +import { ConfigOptions } from "./types"; import { - AuthOptions, - ClientAuthProtocolAdapter, - IsomorphicFetchClientAuthProtocolAdapter + AuthOptions, + ClientAuthProtocolAdapter, + IsomorphicFetchClientAuthProtocolAdapter, } from "./auth"; -const DEFAULT_TENANT = "default_tenant" -const DEFAULT_DATABASE = "default_database" +const DEFAULT_TENANT = "default_tenant"; +const DEFAULT_DATABASE = "default_database"; // interface for tenant interface Tenant { - name: string, + name: string; } // interface for tenant interface Database { - name: string, + name: string; } export class AdminClient { - /** - * @ignore - */ - private api: DefaultApi & ConfigOptions; - private apiAdapter: ClientAuthProtocolAdapter|undefined; - public tenant: string = DEFAULT_TENANT; - public database: string = DEFAULT_DATABASE; + /** + * @ignore + */ + private api: DefaultApi & ConfigOptions; + private apiAdapter: ClientAuthProtocolAdapter | undefined; + public tenant: string = DEFAULT_TENANT; + public database: string = DEFAULT_DATABASE; - /** - * Creates a new AdminClient instance. - * @param {Object} params - The parameters for creating a new client - * @param {string} [params.path] - The base path for the Chroma API. - * @returns {AdminClient} A new AdminClient instance. - * - * @example - * ```typescript - * const client = new AdminClient({ - * path: "http://localhost:8000" - * }); - * ``` - */ - constructor({ - path, - fetchOptions, - auth, - tenant = DEFAULT_TENANT, - database = DEFAULT_DATABASE - }: { - path?: string, - fetchOptions?: RequestInit, - auth?: AuthOptions, - tenant?: string, - database?: string, - } = {}) { - if (path === undefined) path = "http://localhost:8000"; - this.tenant = tenant; - this.database = database; - - const apiConfig: Configuration = new Configuration({ - basePath: path, - }); - if (auth !== undefined) { - this.apiAdapter = new IsomorphicFetchClientAuthProtocolAdapter(new DefaultApi(apiConfig), auth); - this.api = this.apiAdapter.getApi(); - } else { - this.api = new DefaultApi(apiConfig); - } + /** + * Creates a new AdminClient instance. + * @param {Object} params - The parameters for creating a new client + * @param {string} [params.path] - The base path for the Chroma API. + * @returns {AdminClient} A new AdminClient instance. + * + * @example + * ```typescript + * const client = new AdminClient({ + * path: "http://localhost:8000" + * }); + * ``` + */ + constructor({ + path, + fetchOptions, + auth, + tenant = DEFAULT_TENANT, + database = DEFAULT_DATABASE, + }: { + path?: string; + fetchOptions?: RequestInit; + auth?: AuthOptions; + tenant?: string; + database?: string; + } = {}) { + if (path === undefined) path = "http://localhost:8000"; + this.tenant = tenant; + this.database = database; - this.api.options = fetchOptions ?? {}; + const apiConfig: Configuration = new Configuration({ + basePath: path, + }); + if (auth !== undefined) { + this.apiAdapter = new IsomorphicFetchClientAuthProtocolAdapter( + new DefaultApi(apiConfig), + auth, + ); + this.api = this.apiAdapter.getApi(); + } else { + this.api = new DefaultApi(apiConfig); } - /** - * Sets the tenant and database for the client. - * - * @param {Object} params - The parameters for setting tenant and database. - * @param {string} params.tenant - The name of the tenant. - * @param {string} params.database - The name of the database. - * - * @returns {Promise} A promise that returns nothing - * @throws {Error} Any issues - * - * @example - * ```typescript - * await adminClient.setTenant({ - * tenant: "my_tenant", - * database: "my_database", - * }); - * ``` - */ - public async setTenant({ - tenant = DEFAULT_TENANT, - database = DEFAULT_DATABASE - }: { - tenant: string, - database?: string, - }): Promise { - await validateTenantDatabase(this, tenant, database); - this.tenant = tenant; - this.database = database; - } + this.api.options = fetchOptions ?? {}; + } - /** - * Sets the database for the client. - * - * @param {Object} params - The parameters for setting the database. - * @param {string} params.database - The name of the database. - * - * @returns {Promise} A promise that returns nothing - * @throws {Error} Any issues - * - * @example - * ```typescript - * await adminClient.setDatabase({ - * database: "my_database", - * }); - * ``` - */ - public async setDatabase({ - database = DEFAULT_DATABASE - }: { - database?: string, - }): Promise { - await validateTenantDatabase(this, this.tenant, database); - this.database = database; - } + /** + * Sets the tenant and database for the client. + * + * @param {Object} params - The parameters for setting tenant and database. + * @param {string} params.tenant - The name of the tenant. + * @param {string} params.database - The name of the database. + * + * @returns {Promise} A promise that returns nothing + * @throws {Error} Any issues + * + * @example + * ```typescript + * await adminClient.setTenant({ + * tenant: "my_tenant", + * database: "my_database", + * }); + * ``` + */ + public async setTenant({ + tenant = DEFAULT_TENANT, + database = DEFAULT_DATABASE, + }: { + tenant: string; + database?: string; + }): Promise { + await validateTenantDatabase(this, tenant, database); + this.tenant = tenant; + this.database = database; + } - /** - * Creates a new tenant with the specified properties. - * - * @param {Object} params - The parameters for creating a new tenant. - * @param {string} params.name - The name of the tenant. - * - * @returns {Promise} A promise that resolves to the created tenant. - * @throws {Error} If there is an issue creating the tenant. - * - * @example - * ```typescript - * await adminClient.createTenant({ - * name: "my_tenant", - * }); - * ``` - */ - public async createTenant({ - name, - }: { - name: string, - }): Promise { - const newTenant = await this.api - .createTenant({name}, this.api.options) - .then(handleSuccess) - .catch(handleError); + /** + * Sets the database for the client. + * + * @param {Object} params - The parameters for setting the database. + * @param {string} params.database - The name of the database. + * + * @returns {Promise} A promise that returns nothing + * @throws {Error} Any issues + * + * @example + * ```typescript + * await adminClient.setDatabase({ + * database: "my_database", + * }); + * ``` + */ + public async setDatabase({ + database = DEFAULT_DATABASE, + }: { + database?: string; + }): Promise { + await validateTenantDatabase(this, this.tenant, database); + this.database = database; + } - // newTenant is null if successful - if (newTenant && newTenant.error) { - throw new Error(newTenant.error); - } + /** + * Creates a new tenant with the specified properties. + * + * @param {Object} params - The parameters for creating a new tenant. + * @param {string} params.name - The name of the tenant. + * + * @returns {Promise} A promise that resolves to the created tenant. + * @throws {Error} If there is an issue creating the tenant. + * + * @example + * ```typescript + * await adminClient.createTenant({ + * name: "my_tenant", + * }); + * ``` + */ + public async createTenant({ name }: { name: string }): Promise { + const newTenant = await this.api + .createTenant({ name }, this.api.options) + .then(handleSuccess) + .catch(handleError); - return {name: name} as Tenant + // newTenant is null if successful + if (newTenant && newTenant.error) { + throw new Error(newTenant.error); } - /** - * Gets a tenant with the specified properties. - * - * @param {Object} params - The parameters for getting a tenant. - * @param {string} params.name - The name of the tenant. - * - * @returns {Promise} A promise that resolves to the tenant. - * @throws {Error} If there is an issue getting the tenant. - * - * @example - * ```typescript - * await adminClient.getTenant({ - * name: "my_tenant", - * }); - * ``` - */ - public async getTenant({ - name, - }: { - name: string, - }): Promise { - const getTenant = await this.api - .getTenant(name, this.api.options) - .then(handleSuccess) - .catch(handleError); + return { name: name } as Tenant; + } - if (getTenant.error) { - throw new Error(getTenant.error); - } + /** + * Gets a tenant with the specified properties. + * + * @param {Object} params - The parameters for getting a tenant. + * @param {string} params.name - The name of the tenant. + * + * @returns {Promise} A promise that resolves to the tenant. + * @throws {Error} If there is an issue getting the tenant. + * + * @example + * ```typescript + * await adminClient.getTenant({ + * name: "my_tenant", + * }); + * ``` + */ + public async getTenant({ name }: { name: string }): Promise { + const getTenant = await this.api + .getTenant(name, this.api.options) + .then(handleSuccess) + .catch(handleError); - return {name: getTenant.name} as Tenant + if (getTenant.error) { + throw new Error(getTenant.error); } - /** - * Creates a new database with the specified properties. - * - * @param {Object} params - The parameters for creating a new database. - * @param {string} params.name - The name of the database. - * @param {string} params.tenantName - The name of the tenant. - * - * @returns {Promise} A promise that resolves to the created database. - * @throws {Error} If there is an issue creating the database. - * - * @example - * ```typescript - * await adminClient.createDatabase({ - * name: "my_database", - * tenantName: "my_tenant", - * }); - * ``` - */ - public async createDatabase({ - name, - tenantName - }: { - name: string, - tenantName: string, - }): Promise { - const newDatabase = await this.api - .createDatabase(tenantName, {name}, this.api.options) - .then(handleSuccess) - .catch(handleError); + return { name: getTenant.name } as Tenant; + } - // newDatabase is null if successful - if (newDatabase && newDatabase.error) { - throw new Error(newDatabase.error); - } + /** + * Creates a new database with the specified properties. + * + * @param {Object} params - The parameters for creating a new database. + * @param {string} params.name - The name of the database. + * @param {string} params.tenantName - The name of the tenant. + * + * @returns {Promise} A promise that resolves to the created database. + * @throws {Error} If there is an issue creating the database. + * + * @example + * ```typescript + * await adminClient.createDatabase({ + * name: "my_database", + * tenantName: "my_tenant", + * }); + * ``` + */ + public async createDatabase({ + name, + tenantName, + }: { + name: string; + tenantName: string; + }): Promise { + const newDatabase = await this.api + .createDatabase(tenantName, { name }, this.api.options) + .then(handleSuccess) + .catch(handleError); - return {name: name} as Database + // newDatabase is null if successful + if (newDatabase && newDatabase.error) { + throw new Error(newDatabase.error); } - /** - * Gets a database with the specified properties. - * - * @param {Object} params - The parameters for getting a database. - * @param {string} params.name - The name of the database. - * @param {string} params.tenantName - The name of the tenant. - * - * @returns {Promise} A promise that resolves to the database. - * @throws {Error} If there is an issue getting the database. - * - * @example - * ```typescript - * await adminClient.getDatabase({ - * name: "my_database", - * tenantName: "my_tenant", - * }); - * ``` - */ - public async getDatabase({ - name, - tenantName - }: { - name: string, - tenantName: string, - }): Promise { - const getDatabase = await this.api - .getDatabase(name, tenantName, this.api.options) - .then(handleSuccess) - .catch(handleError); + return { name: name } as Database; + } - if (getDatabase.error) { - throw new Error(getDatabase.error); - } + /** + * Gets a database with the specified properties. + * + * @param {Object} params - The parameters for getting a database. + * @param {string} params.name - The name of the database. + * @param {string} params.tenantName - The name of the tenant. + * + * @returns {Promise} A promise that resolves to the database. + * @throws {Error} If there is an issue getting the database. + * + * @example + * ```typescript + * await adminClient.getDatabase({ + * name: "my_database", + * tenantName: "my_tenant", + * }); + * ``` + */ + public async getDatabase({ + name, + tenantName, + }: { + name: string; + tenantName: string; + }): Promise { + const getDatabase = await this.api + .getDatabase(name, tenantName, this.api.options) + .then(handleSuccess) + .catch(handleError); - return {name: getDatabase.name} as Database + if (getDatabase.error) { + throw new Error(getDatabase.error); } + return { name: getDatabase.name } as Database; + } } diff --git a/clients/js/src/ChromaClient.ts b/clients/js/src/ChromaClient.ts index 76edd4e960e..0bc769b6ef6 100644 --- a/clients/js/src/ChromaClient.ts +++ b/clients/js/src/ChromaClient.ts @@ -1,327 +1,357 @@ -import { IEmbeddingFunction } from './embeddings/IEmbeddingFunction'; +import { IEmbeddingFunction } from "./embeddings/IEmbeddingFunction"; import { Configuration, ApiApi as DefaultApi } from "./generated"; import { handleSuccess, handleError } from "./utils"; -import { Collection } from './Collection'; -import { ChromaClientParams, CollectionMetadata, CollectionType, ConfigOptions, CreateCollectionParams, DeleteCollectionParams, GetCollectionParams, GetOrCreateCollectionParams, ListCollectionsParams } from './types'; +import { Collection } from "./Collection"; import { - AuthOptions, - ClientAuthProtocolAdapter, - IsomorphicFetchClientAuthProtocolAdapter + ChromaClientParams, + CollectionMetadata, + CollectionType, + ConfigOptions, + CreateCollectionParams, + DeleteCollectionParams, + GetCollectionParams, + GetOrCreateCollectionParams, + ListCollectionsParams, +} from "./types"; +import { + AuthOptions, + ClientAuthProtocolAdapter, + IsomorphicFetchClientAuthProtocolAdapter, } from "./auth"; -import { DefaultEmbeddingFunction } from './embeddings/DefaultEmbeddingFunction'; -import { AdminClient } from './AdminClient'; +import { DefaultEmbeddingFunction } from "./embeddings/DefaultEmbeddingFunction"; +import { AdminClient } from "./AdminClient"; -const DEFAULT_TENANT = "default_tenant" -const DEFAULT_DATABASE = "default_database" +const DEFAULT_TENANT = "default_tenant"; +const DEFAULT_DATABASE = "default_database"; export class ChromaClient { - /** - * @ignore - */ - private api: DefaultApi & ConfigOptions; - private apiAdapter: ClientAuthProtocolAdapter|undefined; - private tenant: string = DEFAULT_TENANT; - private database: string = DEFAULT_DATABASE; - private _adminClient?: AdminClient - - /** - * Creates a new ChromaClient instance. - * @param {Object} params - The parameters for creating a new client - * @param {string} [params.path] - The base path for the Chroma API. - * @returns {ChromaClient} A new ChromaClient instance. - * - * @example - * ```typescript - * const client = new ChromaClient({ - * path: "http://localhost:8000" - * }); - * ``` - */ - constructor({ - path, - fetchOptions, - auth, - tenant = DEFAULT_TENANT, - database = DEFAULT_DATABASE, - }: ChromaClientParams = {}) { - if (path === undefined) path = "http://localhost:8000"; - this.tenant = tenant; - this.database = database; - - const apiConfig: Configuration = new Configuration({ - basePath: path, - }); + /** + * @ignore + */ + private api: DefaultApi & ConfigOptions; + private apiAdapter: ClientAuthProtocolAdapter | undefined; + private tenant: string = DEFAULT_TENANT; + private database: string = DEFAULT_DATABASE; + private _adminClient?: AdminClient; - if (auth !== undefined) { - this.apiAdapter = new IsomorphicFetchClientAuthProtocolAdapter(new DefaultApi(apiConfig), auth); - this.api = this.apiAdapter.getApi(); - } else { - this.api = new DefaultApi(apiConfig); - } + /** + * Creates a new ChromaClient instance. + * @param {Object} params - The parameters for creating a new client + * @param {string} [params.path] - The base path for the Chroma API. + * @returns {ChromaClient} A new ChromaClient instance. + * + * @example + * ```typescript + * const client = new ChromaClient({ + * path: "http://localhost:8000" + * }); + * ``` + */ + constructor({ + path, + fetchOptions, + auth, + tenant = DEFAULT_TENANT, + database = DEFAULT_DATABASE, + }: ChromaClientParams = {}) { + if (path === undefined) path = "http://localhost:8000"; + this.tenant = tenant; + this.database = database; - this._adminClient = new AdminClient({ - path: path, - fetchOptions: fetchOptions, - auth: auth, - tenant: tenant, - database: database - }); + const apiConfig: Configuration = new Configuration({ + basePath: path, + }); - // TODO: Validate tenant and database on client creation - // this got tricky because: - // - the constructor is sync but the generated api is async - // - we need to inject auth information so a simple rewrite/fetch does not work - - this.api.options = fetchOptions ?? {}; + if (auth !== undefined) { + this.apiAdapter = new IsomorphicFetchClientAuthProtocolAdapter( + new DefaultApi(apiConfig), + auth, + ); + this.api = this.apiAdapter.getApi(); + } else { + this.api = new DefaultApi(apiConfig); } - /** - * Resets the state of the object by making an API call to the reset endpoint. - * - * @returns {Promise} A promise that resolves when the reset operation is complete. - * @throws {Error} If there is an issue resetting the state. - * - * @example - * ```typescript - * await client.reset(); - * ``` - */ - public async reset(): Promise { - return await this.api.reset(this.api.options); - } + this._adminClient = new AdminClient({ + path: path, + fetchOptions: fetchOptions, + auth: auth, + tenant: tenant, + database: database, + }); - /** - * Returns the version of the Chroma API. - * @returns {Promise} A promise that resolves to the version of the Chroma API. - * - * @example - * ```typescript - * const version = await client.version(); - * ``` - */ - public async version(): Promise { - const response = await this.api.version(this.api.options); - return await handleSuccess(response); - } + // TODO: Validate tenant and database on client creation + // this got tricky because: + // - the constructor is sync but the generated api is async + // - we need to inject auth information so a simple rewrite/fetch does not work - /** - * Returns a heartbeat from the Chroma API. - * @returns {Promise} A promise that resolves to the heartbeat from the Chroma API. - * - * @example - * ```typescript - * const heartbeat = await client.heartbeat(); - * ``` - */ - public async heartbeat(): Promise { - const response = await this.api.heartbeat(this.api.options); - let ret = await handleSuccess(response); - return ret["nanosecond heartbeat"] - } + this.api.options = fetchOptions ?? {}; + } - /** - * Creates a new collection with the specified properties. - * - * @param {Object} params - The parameters for creating a new collection. - * @param {string} params.name - The name of the collection. - * @param {CollectionMetadata} [params.metadata] - Optional metadata associated with the collection. - * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. - * - * @returns {Promise} A promise that resolves to the created collection. - * @throws {Error} If there is an issue creating the collection. - * - * @example - * ```typescript - * const collection = await client.createCollection({ - * name: "my_collection", - * metadata: { - * "description": "My first collection" - * } - * }); - * ``` - */ - public async createCollection({ - name, - metadata, - embeddingFunction - }: CreateCollectionParams): Promise { + /** + * Resets the state of the object by making an API call to the reset endpoint. + * + * @returns {Promise} A promise that resolves when the reset operation is complete. + * @throws {Error} If there is an issue resetting the state. + * + * @example + * ```typescript + * await client.reset(); + * ``` + */ + public async reset(): Promise { + return await this.api.reset(this.api.options); + } - if (embeddingFunction === undefined) { - embeddingFunction = new DefaultEmbeddingFunction(); - } - - const newCollection = await this.api - .createCollection(this.tenant, this.database, { - name, - metadata, - }, this.api.options) - .then(handleSuccess) - .catch(handleError); + /** + * Returns the version of the Chroma API. + * @returns {Promise} A promise that resolves to the version of the Chroma API. + * + * @example + * ```typescript + * const version = await client.version(); + * ``` + */ + public async version(): Promise { + const response = await this.api.version(this.api.options); + return await handleSuccess(response); + } - if (newCollection.error) { - throw new Error(newCollection.error); - } + /** + * Returns a heartbeat from the Chroma API. + * @returns {Promise} A promise that resolves to the heartbeat from the Chroma API. + * + * @example + * ```typescript + * const heartbeat = await client.heartbeat(); + * ``` + */ + public async heartbeat(): Promise { + const response = await this.api.heartbeat(this.api.options); + let ret = await handleSuccess(response); + return ret["nanosecond heartbeat"]; + } - return new Collection(name, newCollection.id, this.api, metadata, embeddingFunction); + /** + * Creates a new collection with the specified properties. + * + * @param {Object} params - The parameters for creating a new collection. + * @param {string} params.name - The name of the collection. + * @param {CollectionMetadata} [params.metadata] - Optional metadata associated with the collection. + * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. + * + * @returns {Promise} A promise that resolves to the created collection. + * @throws {Error} If there is an issue creating the collection. + * + * @example + * ```typescript + * const collection = await client.createCollection({ + * name: "my_collection", + * metadata: { + * "description": "My first collection" + * } + * }); + * ``` + */ + public async createCollection({ + name, + metadata, + embeddingFunction, + }: CreateCollectionParams): Promise { + if (embeddingFunction === undefined) { + embeddingFunction = new DefaultEmbeddingFunction(); } - /** - * Gets or creates a collection with the specified properties. - * - * @param {Object} params - The parameters for creating a new collection. - * @param {string} params.name - The name of the collection. - * @param {CollectionMetadata} [params.metadata] - Optional metadata associated with the collection. - * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. - * - * @returns {Promise} A promise that resolves to the got or created collection. - * @throws {Error} If there is an issue getting or creating the collection. - * - * @example - * ```typescript - * const collection = await client.getOrCreateCollection({ - * name: "my_collection", - * metadata: { - * "description": "My first collection" - * } - * }); - * ``` - */ - public async getOrCreateCollection({ - name, - metadata, - embeddingFunction - }: GetOrCreateCollectionParams): Promise { + const newCollection = await this.api + .createCollection( + this.tenant, + this.database, + { + name, + metadata, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); - if (embeddingFunction === undefined) { - embeddingFunction = new DefaultEmbeddingFunction(); - } - - const newCollection = await this.api - .createCollection(this.tenant, this.database, { - name, - metadata, - 'get_or_create': true - }, this.api.options) - .then(handleSuccess) - .catch(handleError); + if (newCollection.error) { + throw new Error(newCollection.error); + } - if (newCollection.error) { - throw new Error(newCollection.error); - } + return new Collection( + name, + newCollection.id, + this.api, + metadata, + embeddingFunction, + ); + } - return new Collection( - name, - newCollection.id, - this.api, - newCollection.metadata, - embeddingFunction - ); + /** + * Gets or creates a collection with the specified properties. + * + * @param {Object} params - The parameters for creating a new collection. + * @param {string} params.name - The name of the collection. + * @param {CollectionMetadata} [params.metadata] - Optional metadata associated with the collection. + * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. + * + * @returns {Promise} A promise that resolves to the got or created collection. + * @throws {Error} If there is an issue getting or creating the collection. + * + * @example + * ```typescript + * const collection = await client.getOrCreateCollection({ + * name: "my_collection", + * metadata: { + * "description": "My first collection" + * } + * }); + * ``` + */ + public async getOrCreateCollection({ + name, + metadata, + embeddingFunction, + }: GetOrCreateCollectionParams): Promise { + if (embeddingFunction === undefined) { + embeddingFunction = new DefaultEmbeddingFunction(); } - /** - * Lists all collections. - * - * @returns {Promise} A promise that resolves to a list of collection names. - * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get. - * @param {PositiveInteger} [params.offset] - Optional offset on the items to get. - * @throws {Error} If there is an issue listing the collections. - * - * @example - * ```typescript - * const collections = await client.listCollections({ - * limit: 10, - * offset: 0, - * }); - * ``` - */ - public async listCollections({ - limit, - offset, - }: ListCollectionsParams = {}): Promise { - const response = await this.api.listCollections( - this.tenant, - this.database, - limit, - offset, - this.api.options); - return handleSuccess(response); - } + const newCollection = await this.api + .createCollection( + this.tenant, + this.database, + { + name, + metadata, + get_or_create: true, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); - /** - * Counts all collections. - * - * @returns {Promise} A promise that resolves to the number of collections. - * @throws {Error} If there is an issue counting the collections. - * - * @example - * ```typescript - * const collections = await client.countCollections(); - * ``` - */ - public async countCollections(): Promise { - const response = await this.api.countCollections(this.tenant, this.database, this.api.options); - return handleSuccess(response); + if (newCollection.error) { + throw new Error(newCollection.error); } - /** - * Gets a collection with the specified name. - * @param {Object} params - The parameters for getting a collection. - * @param {string} params.name - The name of the collection. - * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. - * @returns {Promise} A promise that resolves to the collection. - * @throws {Error} If there is an issue getting the collection. - * - * @example - * ```typescript - * const collection = await client.getCollection({ - * name: "my_collection" - * }); - * ``` - */ - public async getCollection({ - name, - embeddingFunction - }: GetCollectionParams): Promise { - const response = await this.api - .getCollection(name, this.tenant, this.database, this.api.options) - .then(handleSuccess) - .catch(handleError); + return new Collection( + name, + newCollection.id, + this.api, + newCollection.metadata, + embeddingFunction, + ); + } - if (response.error) { - throw new Error(response.error); - } + /** + * Lists all collections. + * + * @returns {Promise} A promise that resolves to a list of collection names. + * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get. + * @param {PositiveInteger} [params.offset] - Optional offset on the items to get. + * @throws {Error} If there is an issue listing the collections. + * + * @example + * ```typescript + * const collections = await client.listCollections({ + * limit: 10, + * offset: 0, + * }); + * ``` + */ + public async listCollections({ + limit, + offset, + }: ListCollectionsParams = {}): Promise { + const response = await this.api.listCollections( + this.tenant, + this.database, + limit, + offset, + this.api.options, + ); + return handleSuccess(response); + } - return new Collection( - response.name, - response.id, - this.api, - response.metadata, - embeddingFunction - ); + /** + * Counts all collections. + * + * @returns {Promise} A promise that resolves to the number of collections. + * @throws {Error} If there is an issue counting the collections. + * + * @example + * ```typescript + * const collections = await client.countCollections(); + * ``` + */ + public async countCollections(): Promise { + const response = await this.api.countCollections( + this.tenant, + this.database, + this.api.options, + ); + return handleSuccess(response); + } - } + /** + * Gets a collection with the specified name. + * @param {Object} params - The parameters for getting a collection. + * @param {string} params.name - The name of the collection. + * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. + * @returns {Promise} A promise that resolves to the collection. + * @throws {Error} If there is an issue getting the collection. + * + * @example + * ```typescript + * const collection = await client.getCollection({ + * name: "my_collection" + * }); + * ``` + */ + public async getCollection({ + name, + embeddingFunction, + }: GetCollectionParams): Promise { + const response = await this.api + .getCollection(name, this.tenant, this.database, this.api.options) + .then(handleSuccess) + .catch(handleError); - /** - * Deletes a collection with the specified name. - * @param {Object} params - The parameters for deleting a collection. - * @param {string} params.name - The name of the collection. - * @returns {Promise} A promise that resolves when the collection is deleted. - * @throws {Error} If there is an issue deleting the collection. - * - * @example - * ```typescript - * await client.deleteCollection({ - * name: "my_collection" - * }); - * ``` - */ - public async deleteCollection({ - name - }: DeleteCollectionParams): Promise { - return await this.api - .deleteCollection(name, this.tenant, this.database, this.api.options) - .then(handleSuccess) - .catch(handleError); + if (response.error) { + throw new Error(response.error); } + return new Collection( + response.name, + response.id, + this.api, + response.metadata, + embeddingFunction, + ); + } + + /** + * Deletes a collection with the specified name. + * @param {Object} params - The parameters for deleting a collection. + * @param {string} params.name - The name of the collection. + * @returns {Promise} A promise that resolves when the collection is deleted. + * @throws {Error} If there is an issue deleting the collection. + * + * @example + * ```typescript + * await client.deleteCollection({ + * name: "my_collection" + * }); + * ``` + */ + public async deleteCollection({ + name, + }: DeleteCollectionParams): Promise { + return await this.api + .deleteCollection(name, this.tenant, this.database, this.api.options) + .then(handleSuccess) + .catch(handleError); + } } diff --git a/clients/js/src/CloudClient.ts b/clients/js/src/CloudClient.ts index 9ce77d2f59d..f2b4f2addf2 100644 --- a/clients/js/src/CloudClient.ts +++ b/clients/js/src/CloudClient.ts @@ -1,46 +1,44 @@ - // create a cloudclient class that takes in an api key and an optional database // this should wrap ChromaClient and specify the auth scheme correctly import { ChromaClient } from "./ChromaClient"; interface CloudClientParams { - apiKey?: string; - database?: string; - cloudHost?: string; - cloudPort?: string; + apiKey?: string; + database?: string; + cloudHost?: string; + cloudPort?: string; } -class CloudClient extends ChromaClient{ - - constructor({apiKey, database, cloudHost, cloudPort}: CloudClientParams) { - // If no API key is provided, try to load it from the environment variable - if (!apiKey) { - apiKey = process.env.CHROMA_API_KEY; - } - if (!apiKey) { - throw new Error("No API key provided"); - } +class CloudClient extends ChromaClient { + constructor({ apiKey, database, cloudHost, cloudPort }: CloudClientParams) { + // If no API key is provided, try to load it from the environment variable + if (!apiKey) { + apiKey = process.env.CHROMA_API_KEY; + } + if (!apiKey) { + throw new Error("No API key provided"); + } - cloudHost = cloudHost || "https://api.trychroma.com"; - cloudPort = cloudPort || "8000"; + cloudHost = cloudHost || "https://api.trychroma.com"; + cloudPort = cloudPort || "8000"; - const path = `${cloudHost}:${cloudPort}`; + const path = `${cloudHost}:${cloudPort}`; - const auth = { - provider: "token", - credentials: apiKey, - providerOptions: { headerType: "X_CHROMA_TOKEN" }, - } + const auth = { + provider: "token", + credentials: apiKey, + providerOptions: { headerType: "X_CHROMA_TOKEN" }, + }; - return new ChromaClient({ - path: path, - auth: auth, - database: database, - }) + return new ChromaClient({ + path: path, + auth: auth, + database: database, + }); - super() - } + super(); + } } export { CloudClient }; diff --git a/clients/js/src/Collection.ts b/clients/js/src/Collection.ts index 82fe3facbd9..430d444bc8a 100644 --- a/clients/js/src/Collection.ts +++ b/clients/js/src/Collection.ts @@ -1,540 +1,551 @@ import { - GetResponse, - QueryResponse, - AddResponse, - CollectionMetadata, - ConfigOptions, - GetParams, - AddParams, - UpsertParams, - ModifyCollectionParams, - UpdateParams, - QueryParams, - PeekParams, - DeleteParams + GetResponse, + QueryResponse, + AddResponse, + CollectionMetadata, + ConfigOptions, + GetParams, + AddParams, + UpsertParams, + ModifyCollectionParams, + UpdateParams, + QueryParams, + PeekParams, + DeleteParams, } from "./types"; -import { IEmbeddingFunction } from './embeddings/IEmbeddingFunction'; +import { IEmbeddingFunction } from "./embeddings/IEmbeddingFunction"; import { ApiApi as DefaultApi } from "./generated"; import { handleError, handleSuccess } from "./utils"; import { toArray, toArrayOfArrays } from "./utils"; - export class Collection { - public name: string; - public id: string; - public metadata: CollectionMetadata | undefined; - /** - * @ignore - */ - private api: DefaultApi & ConfigOptions; - /** - * @ignore - */ - public embeddingFunction: IEmbeddingFunction | undefined; - - /** - * @ignore - */ - constructor( - name: string, - id: string, - api: DefaultApi, - metadata?: CollectionMetadata, - embeddingFunction?: IEmbeddingFunction - ) { - this.name = name; - this.id = id; - this.metadata = metadata; - this.api = api; - if (embeddingFunction !== undefined) - this.embeddingFunction = embeddingFunction; + public name: string; + public id: string; + public metadata: CollectionMetadata | undefined; + /** + * @ignore + */ + private api: DefaultApi & ConfigOptions; + /** + * @ignore + */ + public embeddingFunction: IEmbeddingFunction | undefined; + + /** + * @ignore + */ + constructor( + name: string, + id: string, + api: DefaultApi, + metadata?: CollectionMetadata, + embeddingFunction?: IEmbeddingFunction, + ) { + this.name = name; + this.id = id; + this.metadata = metadata; + this.api = api; + if (embeddingFunction !== undefined) + this.embeddingFunction = embeddingFunction; + } + + /** + * @ignore + */ + private setName(name: string): void { + this.name = name; + } + /** + * @ignore + */ + private setMetadata(metadata: CollectionMetadata | undefined): void { + this.metadata = metadata; + } + + /** + * @ignore + */ + private async validate( + require_embeddings_or_documents: boolean, // set to false in the case of Update + ids: string | string[], + embeddings: number[] | number[][] | undefined, + metadatas?: object | object[], + documents?: string | string[], + ) { + if (require_embeddings_or_documents) { + if (embeddings === undefined && documents === undefined) { + throw new Error("embeddings and documents cannot both be undefined"); + } } - /** - * @ignore - */ - private setName(name: string): void { - this.name = name; - } - /** - * @ignore - */ - private setMetadata(metadata: CollectionMetadata | undefined): void { - this.metadata = metadata; + if (embeddings === undefined && documents !== undefined) { + const documentsArray = toArray(documents); + if (this.embeddingFunction !== undefined) { + embeddings = await this.embeddingFunction.generate(documentsArray); + } else { + throw new Error( + "embeddingFunction is undefined. Please configure an embedding function", + ); + } } + if (embeddings === undefined) + throw new Error("embeddings is undefined but shouldnt be"); - /** - * @ignore - */ - private async validate( - require_embeddings_or_documents: boolean, // set to false in the case of Update - ids: string | string[], - embeddings: number[] | number[][] | undefined, - metadatas?: object | object[], - documents?: string | string[], - ) { + const idsArray = toArray(ids); + const embeddingsArray: number[][] = toArrayOfArrays(embeddings); - if (require_embeddings_or_documents) { - if ((embeddings === undefined) && (documents === undefined)) { - throw new Error( - "embeddings and documents cannot both be undefined", - ); - } - } - - if ((embeddings === undefined) && (documents !== undefined)) { - const documentsArray = toArray(documents); - if (this.embeddingFunction !== undefined) { - embeddings = await this.embeddingFunction.generate(documentsArray); - } else { - throw new Error( - "embeddingFunction is undefined. Please configure an embedding function" - ); - } - } - if (embeddings === undefined) - throw new Error("embeddings is undefined but shouldnt be"); - - const idsArray = toArray(ids); - const embeddingsArray: number[][] = toArrayOfArrays(embeddings); - - let metadatasArray: object[] | undefined; - if (metadatas === undefined) { - metadatasArray = undefined; - } else { - metadatasArray = toArray(metadatas); - } - - let documentsArray: (string | undefined)[] | undefined; - if (documents === undefined) { - documentsArray = undefined; - } else { - documentsArray = toArray(documents); - } - - // validate all ids are strings - for (let i = 0; i < idsArray.length; i += 1) { - if (typeof idsArray[i] !== "string") { - throw new Error( - `Expected ids to be strings, found ${typeof idsArray[i]} at index ${i}` - ); - } - } - - if ( - (embeddingsArray !== undefined && - idsArray.length !== embeddingsArray.length) || - (metadatasArray !== undefined && - idsArray.length !== metadatasArray.length) || - (documentsArray !== undefined && - idsArray.length !== documentsArray.length) - ) { - throw new Error( - "ids, embeddings, metadatas, and documents must all be the same length" - ); - } - - const uniqueIds = new Set(idsArray); - if (uniqueIds.size !== idsArray.length) { - const duplicateIds = idsArray.filter((item, index) => idsArray.indexOf(item) !== index); - throw new Error( - `Expected IDs to be unique, found duplicates for: ${duplicateIds}`, - ); - } - - return [idsArray, embeddingsArray, metadatasArray, documentsArray] + let metadatasArray: object[] | undefined; + if (metadatas === undefined) { + metadatasArray = undefined; + } else { + metadatasArray = toArray(metadatas); } - /** - * Add items to the collection - * @param {Object} params - The parameters for the query. - * @param {ID | IDs} [params.ids] - IDs of the items to add. - * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add. - * @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add. - * @param {Document | Documents} [params.documents] - Optional documents of the items to add. - * @returns {Promise} - The response from the API. True if successful. - * - * @example - * ```typescript - * const response = await collection.add({ - * ids: ["id1", "id2"], - * embeddings: [[1, 2, 3], [4, 5, 6]], - * metadatas: [{ "key": "value" }, { "key": "value" }], - * documents: ["document1", "document2"] - * }); - * ``` - */ - public async add({ - ids, - embeddings, - metadatas, - documents, - }: AddParams): Promise { - - const [idsArray, embeddingsArray, metadatasArray, documentsArray] = await this.validate( - true, - ids, - embeddings, - metadatas, - documents - ) - - const response = await this.api.add(this.id, - { - // @ts-ignore - ids: idsArray, - embeddings: embeddingsArray as number[][], // We know this is defined because of the validate function - // @ts-ignore - documents: documentsArray, - // @ts-ignore - metadatas: metadatasArray, - }, this.api.options) - .then(handleSuccess) - .catch(handleError); - - return response - } - - /** - * Upsert items to the collection - * @param {Object} params - The parameters for the query. - * @param {ID | IDs} [params.ids] - IDs of the items to add. - * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add. - * @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add. - * @param {Document | Documents} [params.documents] - Optional documents of the items to add. - * @returns {Promise} - The response from the API. True if successful. - * - * @example - * ```typescript - * const response = await collection.upsert({ - * ids: ["id1", "id2"], - * embeddings: [[1, 2, 3], [4, 5, 6]], - * metadatas: [{ "key": "value" }, { "key": "value" }], - * documents: ["document1", "document2"], - * }); - * ``` - */ - public async upsert({ - ids, - embeddings, - metadatas, - documents, - }: UpsertParams): Promise { - const [idsArray, embeddingsArray, metadatasArray, documentsArray] = await this.validate( - true, - ids, - embeddings, - metadatas, - documents - ) - - const response = await this.api.upsert(this.id, - { - //@ts-ignore - ids: idsArray, - embeddings: embeddingsArray as number[][], // We know this is defined because of the validate function - //@ts-ignore - documents: documentsArray, - //@ts-ignore - metadatas: metadatasArray, - }, - this.api.options - ) - .then(handleSuccess) - .catch(handleError); - - return response - + let documentsArray: (string | undefined)[] | undefined; + if (documents === undefined) { + documentsArray = undefined; + } else { + documentsArray = toArray(documents); } - /** - * Count the number of items in the collection - * @returns {Promise} - The response from the API. - * - * @example - * ```typescript - * const response = await collection.count(); - * ``` - */ - public async count(): Promise { - const response = await this.api.count(this.id, this.api.options); - return handleSuccess(response); + // validate all ids are strings + for (let i = 0; i < idsArray.length; i += 1) { + if (typeof idsArray[i] !== "string") { + throw new Error( + `Expected ids to be strings, found ${typeof idsArray[ + i + ]} at index ${i}`, + ); + } } - /** - * Modify the collection name or metadata - * @param {Object} params - The parameters for the query. - * @param {string} [params.name] - Optional new name for the collection. - * @param {CollectionMetadata} [params.metadata] - Optional new metadata for the collection. - * @returns {Promise} - The response from the API. - * - * @example - * ```typescript - * const response = await collection.modify({ - * name: "new name", - * metadata: { "key": "value" }, - * }); - * ``` - */ - public async modify({ - name, - metadata - }: ModifyCollectionParams = {}): Promise { - const response = await this.api - .updateCollection( - this.id, - { - new_name: name, - new_metadata: metadata, - }, - this.api.options - ) - .then(handleSuccess) - .catch(handleError); - - this.setName(name || this.name); - this.setMetadata(metadata || this.metadata); - - return response; - } - - /** - * Get items from the collection - * @param {Object} params - The parameters for the query. - * @param {ID | IDs} [params.ids] - Optional IDs of the items to get. - * @param {Where} [params.where] - Optional where clause to filter items by. - * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get. - * @param {PositiveInteger} [params.offset] - Optional offset on the items to get. - * @param {IncludeEnum[]} [params.include] - Optional list of items to include in the response. - * @param {WhereDocument} [params.whereDocument] - Optional where clause to filter items by. - * @returns {Promise} - The response from the server. - * - * @example - * ```typescript - * const response = await collection.get({ - * ids: ["id1", "id2"], - * where: { "key": "value" }, - * limit: 10, - * offset: 0, - * include: ["embeddings", "metadatas", "documents"], - * whereDocument: { $contains: "value" }, - * }); - * ``` - */ - public async get({ - ids, - where, - limit, - offset, - include, - whereDocument, - }: GetParams = {}): Promise { - let idsArray = undefined; - if (ids !== undefined) idsArray = toArray(ids); - - return await this.api - .aGet(this.id, { - ids: idsArray, - where, - limit, - offset, - //@ts-ignore - include, - where_document: whereDocument, - }, this.api.options) - .then(handleSuccess) - .catch(handleError); - } - - /** - * Update the embeddings, documents, and/or metadatas of existing items - * @param {Object} params - The parameters for the query. - * @param {ID | IDs} [params.ids] - The IDs of the items to update. - * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings to update. - * @param {Metadata | Metadatas} [params.metadatas] - Optional metadatas to update. - * @param {Document | Documents} [params.documents] - Optional documents to update. - * @returns {Promise} - The API Response. True if successful. Else, error. - * - * @example - * ```typescript - * const response = await collection.update({ - * ids: ["id1", "id2"], - * embeddings: [[1, 2, 3], [4, 5, 6]], - * metadatas: [{ "key": "value" }, { "key": "value" }], - * documents: ["new document 1", "new document 2"], - * }); - * ``` - */ - public async update({ - ids, - embeddings, - metadatas, - documents, - }: UpdateParams): Promise { - if ( - embeddings === undefined && - documents === undefined && - metadatas === undefined - ) { - throw new Error( - "embeddings, documents, and metadatas cannot all be undefined" - ); - } else if (embeddings === undefined && documents !== undefined) { - const documentsArray = toArray(documents); - if (this.embeddingFunction !== undefined) { - embeddings = await this.embeddingFunction.generate(documentsArray); - } else { - throw new Error( - "embeddingFunction is undefined. Please configure an embedding function" - ); - } - } - - // backend expects None if metadatas is undefined - if (metadatas !== undefined) metadatas = toArray(metadatas); - if (documents !== undefined) documents = toArray(documents); - - var resp = await this.api - .update( - this.id, - { - ids: toArray(ids), - embeddings: embeddings ? toArrayOfArrays(embeddings) : undefined, - documents: documents, - metadatas: metadatas - }, - this.api.options - ) - .then(handleSuccess) - .catch(handleError); - - return resp; + if ( + (embeddingsArray !== undefined && + idsArray.length !== embeddingsArray.length) || + (metadatasArray !== undefined && + idsArray.length !== metadatasArray.length) || + (documentsArray !== undefined && + idsArray.length !== documentsArray.length) + ) { + throw new Error( + "ids, embeddings, metadatas, and documents must all be the same length", + ); } - /** - * Performs a query on the collection using the specified parameters. - * - * @param {Object} params - The parameters for the query. - * @param {Embedding | Embeddings} [params.queryEmbeddings] - Optional query embeddings to use for the search. - * @param {PositiveInteger} [params.nResults] - Optional number of results to return (default is 10). - * @param {Where} [params.where] - Optional query condition to filter results based on metadata values. - * @param {string | string[]} [params.queryTexts] - Optional query text(s) to search for in the collection. - * @param {WhereDocument} [params.whereDocument] - Optional query condition to filter results based on document content. - * @param {IncludeEnum[]} [params.include] - Optional array of fields to include in the result, such as "metadata" and "document". - * - * @returns {Promise} A promise that resolves to the query results. - * @throws {Error} If there is an issue executing the query. - * @example - * // Query the collection using embeddings - * const results = await collection.query({ - * queryEmbeddings: [[0.1, 0.2, ...], ...], - * nResults: 10, - * where: {"name": {"$eq": "John Doe"}}, - * include: ["metadata", "document"] - * }); - * @example - * ```js - * // Query the collection using query text - * const results = await collection.query({ - * queryTexts: "some text", - * nResults: 10, - * where: {"name": {"$eq": "John Doe"}}, - * include: ["metadata", "document"] - * }); - * ``` - * - */ - public async query({ - queryEmbeddings, - nResults, - where, - queryTexts, - whereDocument, - include, - }: QueryParams): Promise { - if (nResults === undefined) nResults = 10 - if (queryEmbeddings === undefined && queryTexts === undefined) { - throw new Error( - "queryEmbeddings and queryTexts cannot both be undefined" - ); - } else if (queryEmbeddings === undefined && queryTexts !== undefined) { - const queryTextsArray = toArray(queryTexts); - if (this.embeddingFunction !== undefined) { - queryEmbeddings = await this.embeddingFunction.generate(queryTextsArray); - } else { - throw new Error( - "embeddingFunction is undefined. Please configure an embedding function" - ); - } - } - if (queryEmbeddings === undefined) - throw new Error("embeddings is undefined but shouldnt be"); - - const query_embeddingsArray = toArrayOfArrays(queryEmbeddings); - - return await this.api - .getNearestNeighbors(this.id, { - query_embeddings: query_embeddingsArray, - where, - n_results: nResults, - where_document: whereDocument, - //@ts-ignore - include: include, - }, this.api.options) - .then(handleSuccess) - .catch(handleError); + const uniqueIds = new Set(idsArray); + if (uniqueIds.size !== idsArray.length) { + const duplicateIds = idsArray.filter( + (item, index) => idsArray.indexOf(item) !== index, + ); + throw new Error( + `Expected IDs to be unique, found duplicates for: ${duplicateIds}`, + ); } - /** - * Peek inside the collection - * @param {Object} params - The parameters for the query. - * @param {PositiveInteger} [params.limit] - Optional number of results to return (default is 10). - * @returns {Promise} A promise that resolves to the query results. - * @throws {Error} If there is an issue executing the query. - * - * @example - * ```typescript - * const results = await collection.peek({ - * limit: 10 - * }); - * ``` - */ - public async peek({ limit }: PeekParams = {}): Promise { - if (limit === undefined) limit = 10; - const response = await this.api.aGet(this.id, { - limit: limit, - }, this.api.options); - return handleSuccess(response); + return [idsArray, embeddingsArray, metadatasArray, documentsArray]; + } + + /** + * Add items to the collection + * @param {Object} params - The parameters for the query. + * @param {ID | IDs} [params.ids] - IDs of the items to add. + * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add. + * @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add. + * @param {Document | Documents} [params.documents] - Optional documents of the items to add. + * @returns {Promise} - The response from the API. True if successful. + * + * @example + * ```typescript + * const response = await collection.add({ + * ids: ["id1", "id2"], + * embeddings: [[1, 2, 3], [4, 5, 6]], + * metadatas: [{ "key": "value" }, { "key": "value" }], + * documents: ["document1", "document2"] + * }); + * ``` + */ + public async add({ + ids, + embeddings, + metadatas, + documents, + }: AddParams): Promise { + const [idsArray, embeddingsArray, metadatasArray, documentsArray] = + await this.validate(true, ids, embeddings, metadatas, documents); + + const response = await this.api + .add( + this.id, + { + // @ts-ignore + ids: idsArray, + embeddings: embeddingsArray as number[][], // We know this is defined because of the validate function + // @ts-ignore + documents: documentsArray, + // @ts-ignore + metadatas: metadatasArray, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + + return response; + } + + /** + * Upsert items to the collection + * @param {Object} params - The parameters for the query. + * @param {ID | IDs} [params.ids] - IDs of the items to add. + * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add. + * @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add. + * @param {Document | Documents} [params.documents] - Optional documents of the items to add. + * @returns {Promise} - The response from the API. True if successful. + * + * @example + * ```typescript + * const response = await collection.upsert({ + * ids: ["id1", "id2"], + * embeddings: [[1, 2, 3], [4, 5, 6]], + * metadatas: [{ "key": "value" }, { "key": "value" }], + * documents: ["document1", "document2"], + * }); + * ``` + */ + public async upsert({ + ids, + embeddings, + metadatas, + documents, + }: UpsertParams): Promise { + const [idsArray, embeddingsArray, metadatasArray, documentsArray] = + await this.validate(true, ids, embeddings, metadatas, documents); + + const response = await this.api + .upsert( + this.id, + { + //@ts-ignore + ids: idsArray, + embeddings: embeddingsArray as number[][], // We know this is defined because of the validate function + //@ts-ignore + documents: documentsArray, + //@ts-ignore + metadatas: metadatasArray, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + + return response; + } + + /** + * Count the number of items in the collection + * @returns {Promise} - The response from the API. + * + * @example + * ```typescript + * const response = await collection.count(); + * ``` + */ + public async count(): Promise { + const response = await this.api.count(this.id, this.api.options); + return handleSuccess(response); + } + + /** + * Modify the collection name or metadata + * @param {Object} params - The parameters for the query. + * @param {string} [params.name] - Optional new name for the collection. + * @param {CollectionMetadata} [params.metadata] - Optional new metadata for the collection. + * @returns {Promise} - The response from the API. + * + * @example + * ```typescript + * const response = await collection.modify({ + * name: "new name", + * metadata: { "key": "value" }, + * }); + * ``` + */ + public async modify({ + name, + metadata, + }: ModifyCollectionParams = {}): Promise { + const response = await this.api + .updateCollection( + this.id, + { + new_name: name, + new_metadata: metadata, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + + this.setName(name || this.name); + this.setMetadata(metadata || this.metadata); + + return response; + } + + /** + * Get items from the collection + * @param {Object} params - The parameters for the query. + * @param {ID | IDs} [params.ids] - Optional IDs of the items to get. + * @param {Where} [params.where] - Optional where clause to filter items by. + * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get. + * @param {PositiveInteger} [params.offset] - Optional offset on the items to get. + * @param {IncludeEnum[]} [params.include] - Optional list of items to include in the response. + * @param {WhereDocument} [params.whereDocument] - Optional where clause to filter items by. + * @returns {Promise} - The response from the server. + * + * @example + * ```typescript + * const response = await collection.get({ + * ids: ["id1", "id2"], + * where: { "key": "value" }, + * limit: 10, + * offset: 0, + * include: ["embeddings", "metadatas", "documents"], + * whereDocument: { $contains: "value" }, + * }); + * ``` + */ + public async get({ + ids, + where, + limit, + offset, + include, + whereDocument, + }: GetParams = {}): Promise { + let idsArray = undefined; + if (ids !== undefined) idsArray = toArray(ids); + + return await this.api + .aGet( + this.id, + { + ids: idsArray, + where, + limit, + offset, + //@ts-ignore + include, + where_document: whereDocument, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + } + + /** + * Update the embeddings, documents, and/or metadatas of existing items + * @param {Object} params - The parameters for the query. + * @param {ID | IDs} [params.ids] - The IDs of the items to update. + * @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings to update. + * @param {Metadata | Metadatas} [params.metadatas] - Optional metadatas to update. + * @param {Document | Documents} [params.documents] - Optional documents to update. + * @returns {Promise} - The API Response. True if successful. Else, error. + * + * @example + * ```typescript + * const response = await collection.update({ + * ids: ["id1", "id2"], + * embeddings: [[1, 2, 3], [4, 5, 6]], + * metadatas: [{ "key": "value" }, { "key": "value" }], + * documents: ["new document 1", "new document 2"], + * }); + * ``` + */ + public async update({ + ids, + embeddings, + metadatas, + documents, + }: UpdateParams): Promise { + if ( + embeddings === undefined && + documents === undefined && + metadatas === undefined + ) { + throw new Error( + "embeddings, documents, and metadatas cannot all be undefined", + ); + } else if (embeddings === undefined && documents !== undefined) { + const documentsArray = toArray(documents); + if (this.embeddingFunction !== undefined) { + embeddings = await this.embeddingFunction.generate(documentsArray); + } else { + throw new Error( + "embeddingFunction is undefined. Please configure an embedding function", + ); + } } - /** - * Deletes items from the collection. - * @param {Object} params - The parameters for deleting items from the collection. - * @param {ID | IDs} [params.ids] - Optional ID or array of IDs of items to delete. - * @param {Where} [params.where] - Optional query condition to filter items to delete based on metadata values. - * @param {WhereDocument} [params.whereDocument] - Optional query condition to filter items to delete based on document content. - * @returns {Promise} A promise that resolves to the IDs of the deleted items. - * @throws {Error} If there is an issue deleting items from the collection. - * - * @example - * ```typescript - * const results = await collection.delete({ - * ids: "some_id", - * where: {"name": {"$eq": "John Doe"}}, - * whereDocument: {"$contains":"search_string"} - * }); - * ``` - */ - public async delete({ - ids, - where, - whereDocument - }: DeleteParams = {}): Promise { - let idsArray = undefined; - if (ids !== undefined) idsArray = toArray(ids); - return await this.api - .aDelete(this.id, { ids: idsArray, where: where, where_document: whereDocument }, this.api.options) - .then(handleSuccess) - .catch(handleError); + // backend expects None if metadatas is undefined + if (metadatas !== undefined) metadatas = toArray(metadatas); + if (documents !== undefined) documents = toArray(documents); + + var resp = await this.api + .update( + this.id, + { + ids: toArray(ids), + embeddings: embeddings ? toArrayOfArrays(embeddings) : undefined, + documents: documents, + metadatas: metadatas, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + + return resp; + } + + /** + * Performs a query on the collection using the specified parameters. + * + * @param {Object} params - The parameters for the query. + * @param {Embedding | Embeddings} [params.queryEmbeddings] - Optional query embeddings to use for the search. + * @param {PositiveInteger} [params.nResults] - Optional number of results to return (default is 10). + * @param {Where} [params.where] - Optional query condition to filter results based on metadata values. + * @param {string | string[]} [params.queryTexts] - Optional query text(s) to search for in the collection. + * @param {WhereDocument} [params.whereDocument] - Optional query condition to filter results based on document content. + * @param {IncludeEnum[]} [params.include] - Optional array of fields to include in the result, such as "metadata" and "document". + * + * @returns {Promise} A promise that resolves to the query results. + * @throws {Error} If there is an issue executing the query. + * @example + * // Query the collection using embeddings + * const results = await collection.query({ + * queryEmbeddings: [[0.1, 0.2, ...], ...], + * nResults: 10, + * where: {"name": {"$eq": "John Doe"}}, + * include: ["metadata", "document"] + * }); + * @example + * ```js + * // Query the collection using query text + * const results = await collection.query({ + * queryTexts: "some text", + * nResults: 10, + * where: {"name": {"$eq": "John Doe"}}, + * include: ["metadata", "document"] + * }); + * ``` + * + */ + public async query({ + queryEmbeddings, + nResults, + where, + queryTexts, + whereDocument, + include, + }: QueryParams): Promise { + if (nResults === undefined) nResults = 10; + if (queryEmbeddings === undefined && queryTexts === undefined) { + throw new Error( + "queryEmbeddings and queryTexts cannot both be undefined", + ); + } else if (queryEmbeddings === undefined && queryTexts !== undefined) { + const queryTextsArray = toArray(queryTexts); + if (this.embeddingFunction !== undefined) { + queryEmbeddings = + await this.embeddingFunction.generate(queryTextsArray); + } else { + throw new Error( + "embeddingFunction is undefined. Please configure an embedding function", + ); + } } + if (queryEmbeddings === undefined) + throw new Error("embeddings is undefined but shouldnt be"); + + const query_embeddingsArray = toArrayOfArrays(queryEmbeddings); + + return await this.api + .getNearestNeighbors( + this.id, + { + query_embeddings: query_embeddingsArray, + where, + n_results: nResults, + where_document: whereDocument, + //@ts-ignore + include: include, + }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + } + + /** + * Peek inside the collection + * @param {Object} params - The parameters for the query. + * @param {PositiveInteger} [params.limit] - Optional number of results to return (default is 10). + * @returns {Promise} A promise that resolves to the query results. + * @throws {Error} If there is an issue executing the query. + * + * @example + * ```typescript + * const results = await collection.peek({ + * limit: 10 + * }); + * ``` + */ + public async peek({ limit }: PeekParams = {}): Promise { + if (limit === undefined) limit = 10; + const response = await this.api.aGet( + this.id, + { + limit: limit, + }, + this.api.options, + ); + return handleSuccess(response); + } + + /** + * Deletes items from the collection. + * @param {Object} params - The parameters for deleting items from the collection. + * @param {ID | IDs} [params.ids] - Optional ID or array of IDs of items to delete. + * @param {Where} [params.where] - Optional query condition to filter items to delete based on metadata values. + * @param {WhereDocument} [params.whereDocument] - Optional query condition to filter items to delete based on document content. + * @returns {Promise} A promise that resolves to the IDs of the deleted items. + * @throws {Error} If there is an issue deleting items from the collection. + * + * @example + * ```typescript + * const results = await collection.delete({ + * ids: "some_id", + * where: {"name": {"$eq": "John Doe"}}, + * whereDocument: {"$contains":"search_string"} + * }); + * ``` + */ + public async delete({ + ids, + where, + whereDocument, + }: DeleteParams = {}): Promise { + let idsArray = undefined; + if (ids !== undefined) idsArray = toArray(ids); + return await this.api + .aDelete( + this.id, + { ids: idsArray, where: where, where_document: whereDocument }, + this.api.options, + ) + .then(handleSuccess) + .catch(handleError); + } } diff --git a/clients/js/src/auth.ts b/clients/js/src/auth.ts index 4f833f97d61..415c98e14b4 100644 --- a/clients/js/src/auth.ts +++ b/clients/js/src/auth.ts @@ -1,321 +1,376 @@ -import {ApiApi as DefaultApi} from "./generated"; +import { ApiApi as DefaultApi } from "./generated"; export interface ClientAuthProvider { - /** - * Abstract method for authenticating a client. - */ - authenticate(): ClientAuthResponse; + /** + * Abstract method for authenticating a client. + */ + authenticate(): ClientAuthResponse; } export interface ClientAuthConfigurationProvider { - /** - * Abstract method for getting the configuration for the client. - */ - getConfig(): T; + /** + * Abstract method for getting the configuration for the client. + */ + getConfig(): T; } export interface ClientAuthCredentialsProvider { - /** - * Abstract method for getting the credentials for the client. - * @param user - */ - getCredentials(user?: string): T; + /** + * Abstract method for getting the credentials for the client. + * @param user + */ + getCredentials(user?: string): T; } enum AuthInfoType { - COOKIE = "cookie", - HEADER = "header", - URL = "url", - METADATA = "metadata" - + COOKIE = "cookie", + HEADER = "header", + URL = "url", + METADATA = "metadata", } export interface ClientAuthResponse { - getAuthInfoType(): AuthInfoType; + getAuthInfoType(): AuthInfoType; - getAuthInfo(): { key: string, value: string }; + getAuthInfo(): { key: string; value: string }; } - export interface AbstractCredentials { - getCredentials(): T; + getCredentials(): T; } export interface ClientAuthProtocolAdapter { - injectCredentials(injectionContext: T): T; + injectCredentials(injectionContext: T): T; - getApi(): any; + getApi(): any; } - class SecretStr { - constructor(private readonly secret: string) { - } + constructor(private readonly secret: string) {} - getSecret(): string { - return this.secret; - } + getSecret(): string { + return this.secret; + } } const base64Encode = (str: string): string => { - return Buffer.from(str).toString('base64'); + return Buffer.from(str).toString("base64"); }; class BasicAuthCredentials implements AbstractCredentials { - private readonly credentials: SecretStr; + private readonly credentials: SecretStr; - constructor(_creds: string) { - this.credentials = new SecretStr(base64Encode(_creds)) - } + constructor(_creds: string) { + this.credentials = new SecretStr(base64Encode(_creds)); + } - getCredentials(): SecretStr { - //encode base64 - return this.credentials; - } + getCredentials(): SecretStr { + //encode base64 + return this.credentials; + } } - class BasicAuthClientAuthResponse implements ClientAuthResponse { - constructor(private readonly credentials: BasicAuthCredentials) { - } - - getAuthInfo(): { key: string; value: string } { - return {key: "Authorization", value: "Basic " + this.credentials.getCredentials().getSecret()}; - } - - getAuthInfoType(): AuthInfoType { - return AuthInfoType.HEADER; - } + constructor(private readonly credentials: BasicAuthCredentials) {} + + getAuthInfo(): { key: string; value: string } { + return { + key: "Authorization", + value: "Basic " + this.credentials.getCredentials().getSecret(), + }; + } + + getAuthInfoType(): AuthInfoType { + return AuthInfoType.HEADER; + } } -export class BasicAuthCredentialsProvider implements ClientAuthCredentialsProvider { - private readonly credentials: BasicAuthCredentials; - - /** - * Creates a new BasicAuthCredentialsProvider. This provider loads credentials from provided text credentials or from the environment variable CHROMA_CLIENT_AUTH_CREDENTIALS. - * @param _creds - The credentials - * @throws {Error} If neither credentials provider or text credentials are supplied. - */ - - constructor(_creds: string | undefined) { - if (_creds === undefined && !process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) throw new Error("Credentials must be supplied via environment variable (CHROMA_CLIENT_AUTH_CREDENTIALS) or passed in as configuration."); - this.credentials = new BasicAuthCredentials((_creds ?? process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) as string); - } - - getCredentials(): BasicAuthCredentials { - return this.credentials; - } +export class BasicAuthCredentialsProvider + implements ClientAuthCredentialsProvider +{ + private readonly credentials: BasicAuthCredentials; + + /** + * Creates a new BasicAuthCredentialsProvider. This provider loads credentials from provided text credentials or from the environment variable CHROMA_CLIENT_AUTH_CREDENTIALS. + * @param _creds - The credentials + * @throws {Error} If neither credentials provider or text credentials are supplied. + */ + + constructor(_creds: string | undefined) { + if (_creds === undefined && !process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) + throw new Error( + "Credentials must be supplied via environment variable (CHROMA_CLIENT_AUTH_CREDENTIALS) or passed in as configuration.", + ); + this.credentials = new BasicAuthCredentials( + (_creds ?? process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) as string, + ); + } + + getCredentials(): BasicAuthCredentials { + return this.credentials; + } } class BasicAuthClientAuthProvider implements ClientAuthProvider { - private readonly credentialsProvider: ClientAuthCredentialsProvider; - - /** - * Creates a new BasicAuthClientAuthProvider. - * @param options - The options for the authentication provider. - * @param options.textCredentials - The credentials for the authentication provider. - * @param options.credentialsProvider - The credentials provider for the authentication provider. - * @throws {Error} If neither credentials provider or text credentials are supplied. - */ - - constructor(options: { - textCredentials: any; - credentialsProvider: ClientAuthCredentialsProvider | undefined - }) { - if (!options.credentialsProvider && !options.textCredentials) { - throw new Error("Either credentials provider or text credentials must be supplied."); - } - this.credentialsProvider = options.credentialsProvider || new BasicAuthCredentialsProvider(options.textCredentials); - } - - authenticate(): ClientAuthResponse { - return new BasicAuthClientAuthResponse(this.credentialsProvider.getCredentials()); + private readonly credentialsProvider: ClientAuthCredentialsProvider; + + /** + * Creates a new BasicAuthClientAuthProvider. + * @param options - The options for the authentication provider. + * @param options.textCredentials - The credentials for the authentication provider. + * @param options.credentialsProvider - The credentials provider for the authentication provider. + * @throws {Error} If neither credentials provider or text credentials are supplied. + */ + + constructor(options: { + textCredentials: any; + credentialsProvider: ClientAuthCredentialsProvider | undefined; + }) { + if (!options.credentialsProvider && !options.textCredentials) { + throw new Error( + "Either credentials provider or text credentials must be supplied.", + ); } + this.credentialsProvider = + options.credentialsProvider || + new BasicAuthCredentialsProvider(options.textCredentials); + } + + authenticate(): ClientAuthResponse { + return new BasicAuthClientAuthResponse( + this.credentialsProvider.getCredentials(), + ); + } } class TokenAuthCredentials implements AbstractCredentials { - private readonly credentials: SecretStr; + private readonly credentials: SecretStr; - constructor(_creds: string) { - this.credentials = new SecretStr(_creds) - } + constructor(_creds: string) { + this.credentials = new SecretStr(_creds); + } - getCredentials(): SecretStr { - return this.credentials; - } + getCredentials(): SecretStr { + return this.credentials; + } } -export class TokenCredentialsProvider implements ClientAuthCredentialsProvider { - private readonly credentials: TokenAuthCredentials; - - constructor(_creds: string | undefined) { - if (_creds === undefined && !process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) throw new Error("Credentials must be supplied via environment variable (CHROMA_CLIENT_AUTH_CREDENTIALS) or passed in as configuration."); - this.credentials = new TokenAuthCredentials((_creds ?? process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) as string); - } - - getCredentials(): TokenAuthCredentials { - return this.credentials; - } +export class TokenCredentialsProvider + implements ClientAuthCredentialsProvider +{ + private readonly credentials: TokenAuthCredentials; + + constructor(_creds: string | undefined) { + if (_creds === undefined && !process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) + throw new Error( + "Credentials must be supplied via environment variable (CHROMA_CLIENT_AUTH_CREDENTIALS) or passed in as configuration.", + ); + this.credentials = new TokenAuthCredentials( + (_creds ?? process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) as string, + ); + } + + getCredentials(): TokenAuthCredentials { + return this.credentials; + } } export class TokenClientAuthProvider implements ClientAuthProvider { - private readonly credentialsProvider: ClientAuthCredentialsProvider; - private readonly providerOptions: { headerType: TokenHeaderType }; - - constructor(options: { - textCredentials: any; - credentialsProvider: ClientAuthCredentialsProvider | undefined, - providerOptions?: { headerType: TokenHeaderType } - }) { - if (!options.credentialsProvider && !options.textCredentials) { - throw new Error("Either credentials provider or text credentials must be supplied."); - } - if (options.providerOptions === undefined || !options.providerOptions.hasOwnProperty("headerType")) { - this.providerOptions = {headerType: "AUTHORIZATION"}; - } else { - this.providerOptions = {headerType: options.providerOptions.headerType}; - } - this.credentialsProvider = options.credentialsProvider || new TokenCredentialsProvider(options.textCredentials); + private readonly credentialsProvider: ClientAuthCredentialsProvider; + private readonly providerOptions: { headerType: TokenHeaderType }; + + constructor(options: { + textCredentials: any; + credentialsProvider: ClientAuthCredentialsProvider | undefined; + providerOptions?: { headerType: TokenHeaderType }; + }) { + if (!options.credentialsProvider && !options.textCredentials) { + throw new Error( + "Either credentials provider or text credentials must be supplied.", + ); } - - authenticate(): ClientAuthResponse { - return new TokenClientAuthResponse(this.credentialsProvider.getCredentials(), this.providerOptions.headerType); + if ( + options.providerOptions === undefined || + !options.providerOptions.hasOwnProperty("headerType") + ) { + this.providerOptions = { headerType: "AUTHORIZATION" }; + } else { + this.providerOptions = { headerType: options.providerOptions.headerType }; } - + this.credentialsProvider = + options.credentialsProvider || + new TokenCredentialsProvider(options.textCredentials); + } + + authenticate(): ClientAuthResponse { + return new TokenClientAuthResponse( + this.credentialsProvider.getCredentials(), + this.providerOptions.headerType, + ); + } } - -type TokenHeaderType = 'AUTHORIZATION' | 'X_CHROMA_TOKEN'; - -const TokenHeader: Record { key: string; value: string; }> = { - AUTHORIZATION: (value: string) => ({key: "Authorization", value: `Bearer ${value}`}), - X_CHROMA_TOKEN: (value: string) => ({key: "X-Chroma-Token", value: value}) -} +type TokenHeaderType = "AUTHORIZATION" | "X_CHROMA_TOKEN"; + +const TokenHeader: Record< + TokenHeaderType, + (value: string) => { key: string; value: string } +> = { + AUTHORIZATION: (value: string) => ({ + key: "Authorization", + value: `Bearer ${value}`, + }), + X_CHROMA_TOKEN: (value: string) => ({ key: "X-Chroma-Token", value: value }), +}; class TokenClientAuthResponse implements ClientAuthResponse { - constructor(private readonly credentials: TokenAuthCredentials, private readonly headerType: TokenHeaderType = 'AUTHORIZATION') { - } - - getAuthInfo(): { key: string; value: string } { - if (this.headerType === 'AUTHORIZATION') { - return TokenHeader.AUTHORIZATION(this.credentials.getCredentials().getSecret()); - } else if (this.headerType === 'X_CHROMA_TOKEN') { - return TokenHeader.X_CHROMA_TOKEN(this.credentials.getCredentials().getSecret()); - } else { - throw new Error("Invalid header type: " + this.headerType + ". Valid types are: " + Object.keys(TokenHeader).join(", ")); - } + constructor( + private readonly credentials: TokenAuthCredentials, + private readonly headerType: TokenHeaderType = "AUTHORIZATION", + ) {} + + getAuthInfo(): { key: string; value: string } { + if (this.headerType === "AUTHORIZATION") { + return TokenHeader.AUTHORIZATION( + this.credentials.getCredentials().getSecret(), + ); + } else if (this.headerType === "X_CHROMA_TOKEN") { + return TokenHeader.X_CHROMA_TOKEN( + this.credentials.getCredentials().getSecret(), + ); + } else { + throw new Error( + "Invalid header type: " + + this.headerType + + ". Valid types are: " + + Object.keys(TokenHeader).join(", "), + ); } + } - getAuthInfoType(): AuthInfoType { - return AuthInfoType.HEADER; - } + getAuthInfoType(): AuthInfoType { + return AuthInfoType.HEADER; + } } - -export class IsomorphicFetchClientAuthProtocolAdapter implements ClientAuthProtocolAdapter { - authProvider: ClientAuthProvider | undefined; - wrapperApi: DefaultApi | undefined; - - /** - * Creates a new adapter of IsomorphicFetchClientAuthProtocolAdapter. - * @param api - The API to wrap. - * @param authConfiguration - The configuration for the authentication provider. - */ - - constructor(private api: DefaultApi, authConfiguration: AuthOptions) { - - switch (authConfiguration.provider) { - case "basic": - this.authProvider = new BasicAuthClientAuthProvider({ - textCredentials: authConfiguration.credentials, - credentialsProvider: authConfiguration.credentialsProvider - }); - break; - case "token": - this.authProvider = new TokenClientAuthProvider({ - textCredentials: authConfiguration.credentials, - credentialsProvider: authConfiguration.credentialsProvider, - providerOptions: authConfiguration.providerOptions - }); - break; - default: - this.authProvider = undefined; - break; - } - if (this.authProvider !== undefined) { - this.wrapperApi = this.wrapMethods(this.api); - } - } - - getApi(): DefaultApi { - return this.wrapperApi ?? this.api; +export class IsomorphicFetchClientAuthProtocolAdapter + implements ClientAuthProtocolAdapter +{ + authProvider: ClientAuthProvider | undefined; + wrapperApi: DefaultApi | undefined; + + /** + * Creates a new adapter of IsomorphicFetchClientAuthProtocolAdapter. + * @param api - The API to wrap. + * @param authConfiguration - The configuration for the authentication provider. + */ + + constructor( + private api: DefaultApi, + authConfiguration: AuthOptions, + ) { + switch (authConfiguration.provider) { + case "basic": + this.authProvider = new BasicAuthClientAuthProvider({ + textCredentials: authConfiguration.credentials, + credentialsProvider: authConfiguration.credentialsProvider, + }); + break; + case "token": + this.authProvider = new TokenClientAuthProvider({ + textCredentials: authConfiguration.credentials, + credentialsProvider: authConfiguration.credentialsProvider, + providerOptions: authConfiguration.providerOptions, + }); + break; + default: + this.authProvider = undefined; + break; } - - getAllMethods(obj: any): string[] { - let methods: string[] = []; - let currentObj = obj; - - do { - const objMethods = Object.getOwnPropertyNames(currentObj) - .filter(name => typeof currentObj[name] === 'function' && name !== 'constructor'); - - methods = methods.concat(objMethods); - currentObj = Object.getPrototypeOf(currentObj); - } while (currentObj); - - return methods; + if (this.authProvider !== undefined) { + this.wrapperApi = this.wrapMethods(this.api); } - - wrapMethods(obj: any): any { - let self = this; - const methodNames = Object.getOwnPropertyNames(Object.getPrototypeOf(obj)) - .filter(name => typeof obj[name] === 'function' && name !== 'constructor'); - - return new Proxy(obj, { - get(target, prop: string) { - if (methodNames.includes(prop)) { - return new Proxy(target[prop], { - apply(fn, thisArg, args) { - const modifiedArgs = args.map(arg => { - if (arg && typeof arg === 'object' && 'method' in arg) { - return self.injectCredentials(arg as RequestInit); - } - return arg; - }); - if (Object.keys(modifiedArgs[modifiedArgs.length - 1]).length === 0) { - modifiedArgs[modifiedArgs.length - 1] = self.injectCredentials({} as RequestInit); - } else { - modifiedArgs[modifiedArgs.length - 1] = self.injectCredentials(modifiedArgs[modifiedArgs.length - 1] as RequestInit); - } - return fn.apply(thisArg, modifiedArgs); - } - }); + } + + getApi(): DefaultApi { + return this.wrapperApi ?? this.api; + } + + getAllMethods(obj: any): string[] { + let methods: string[] = []; + let currentObj = obj; + + do { + const objMethods = Object.getOwnPropertyNames(currentObj).filter( + (name) => + typeof currentObj[name] === "function" && name !== "constructor", + ); + + methods = methods.concat(objMethods); + currentObj = Object.getPrototypeOf(currentObj); + } while (currentObj); + + return methods; + } + + wrapMethods(obj: any): any { + let self = this; + const methodNames = Object.getOwnPropertyNames( + Object.getPrototypeOf(obj), + ).filter( + (name) => typeof obj[name] === "function" && name !== "constructor", + ); + + return new Proxy(obj, { + get(target, prop: string) { + if (methodNames.includes(prop)) { + return new Proxy(target[prop], { + apply(fn, thisArg, args) { + const modifiedArgs = args.map((arg) => { + if (arg && typeof arg === "object" && "method" in arg) { + return self.injectCredentials(arg as RequestInit); } - return target[prop]; - } - }); - } - - injectCredentials(injectionContext: RequestInit): RequestInit { - const authInfo = this.authProvider?.authenticate().getAuthInfo(); - if (authInfo) { - const {key, value} = authInfo; - injectionContext = { - ...injectionContext, - headers: { - [key]: value - }, - } + return arg; + }); + if ( + Object.keys(modifiedArgs[modifiedArgs.length - 1]).length === 0 + ) { + modifiedArgs[modifiedArgs.length - 1] = self.injectCredentials( + {} as RequestInit, + ); + } else { + modifiedArgs[modifiedArgs.length - 1] = self.injectCredentials( + modifiedArgs[modifiedArgs.length - 1] as RequestInit, + ); + } + return fn.apply(thisArg, modifiedArgs); + }, + }); } - return injectionContext; + return target[prop]; + }, + }); + } + + injectCredentials(injectionContext: RequestInit): RequestInit { + const authInfo = this.authProvider?.authenticate().getAuthInfo(); + if (authInfo) { + const { key, value } = authInfo; + injectionContext = { + ...injectionContext, + headers: { + [key]: value, + }, + }; } + return injectionContext; + } } - export type AuthOptions = { - provider: ClientAuthProvider | string | undefined, - credentialsProvider?: ClientAuthCredentialsProvider | undefined, - configProvider?: ClientAuthConfigurationProvider | undefined, - credentials?: any | undefined, - providerOptions?: any | undefined -} + provider: ClientAuthProvider | string | undefined; + credentialsProvider?: ClientAuthCredentialsProvider | undefined; + configProvider?: ClientAuthConfigurationProvider | undefined; + credentials?: any | undefined; + providerOptions?: any | undefined; +}; diff --git a/clients/js/src/embeddings/DefaultEmbeddingFunction.ts b/clients/js/src/embeddings/DefaultEmbeddingFunction.ts index 6ced79bbd48..be2394a7d22 100644 --- a/clients/js/src/embeddings/DefaultEmbeddingFunction.ts +++ b/clients/js/src/embeddings/DefaultEmbeddingFunction.ts @@ -39,21 +39,21 @@ export class DefaultEmbeddingFunction implements IEmbeddingFunction { public async generate(texts: string[]): Promise { await this.loadClient(); - // Store a promise that resolves to the pipeline + // Store a promise that resolves to the pipeline this.pipelinePromise = new Promise(async (resolve, reject) => { try { - const pipeline = this.transformersApi + const pipeline = this.transformersApi; - const quantized = this.quantized - const revision = this.revision - const progress_callback = this.progress_callback + const quantized = this.quantized; + const revision = this.revision; + const progress_callback = this.progress_callback; resolve( await pipeline("feature-extraction", this.model, { quantized, revision, progress_callback, - }) + }), ); } catch (e) { reject(e); @@ -66,34 +66,36 @@ export class DefaultEmbeddingFunction implements IEmbeddingFunction { } private async loadClient() { - if(this.transformersApi) return; - try { - // eslint-disable-next-line global-require,import/no-extraneous-dependencies - let { pipeline } = await DefaultEmbeddingFunction.import(); - TransformersApi = pipeline; - } catch (_a) { - // @ts-ignore - if (_a.code === 'MODULE_NOT_FOUND') { - throw new Error("Please install the chromadb-default-embed package to use the DefaultEmbeddingFunction, `npm install -S chromadb-default-embed`"); - } - throw _a; // Re-throw other errors + if (this.transformersApi) return; + try { + // eslint-disable-next-line global-require,import/no-extraneous-dependencies + let { pipeline } = await DefaultEmbeddingFunction.import(); + TransformersApi = pipeline; + } catch (_a) { + // @ts-ignore + if (_a.code === "MODULE_NOT_FOUND") { + throw new Error( + "Please install the chromadb-default-embed package to use the DefaultEmbeddingFunction, `npm install -S chromadb-default-embed`", + ); } - this.transformersApi = TransformersApi; + throw _a; // Re-throw other errors + } + this.transformersApi = TransformersApi; } /** @ignore */ static async import(): Promise<{ - // @ts-ignore - pipeline: typeof import("chromadb-default-embed"); + // @ts-ignore + pipeline: typeof import("chromadb-default-embed"); }> { - try { - // @ts-ignore - const { pipeline } = await import("chromadb-default-embed"); - return { pipeline }; - } catch (e) { - throw new Error( - "Please install chromadb-default-embed as a dependency with, e.g. `yarn add chromadb-default-embed`" - ); - } + try { + // @ts-ignore + const { pipeline } = await import("chromadb-default-embed"); + return { pipeline }; + } catch (e) { + throw new Error( + "Please install chromadb-default-embed as a dependency with, e.g. `yarn add chromadb-default-embed`", + ); + } } } diff --git a/clients/js/src/embeddings/GoogleGeminiEmbeddingFunction.ts b/clients/js/src/embeddings/GoogleGeminiEmbeddingFunction.ts index a1ab2abe995..54a356481b0 100644 --- a/clients/js/src/embeddings/GoogleGeminiEmbeddingFunction.ts +++ b/clients/js/src/embeddings/GoogleGeminiEmbeddingFunction.ts @@ -3,67 +3,76 @@ import { IEmbeddingFunction } from "./IEmbeddingFunction"; let googleGenAiApi: any; export class GoogleGenerativeAiEmbeddingFunction implements IEmbeddingFunction { - private api_key: string; - private model: string; - private googleGenAiApi?: any; - private taskType: string; + private api_key: string; + private model: string; + private googleGenAiApi?: any; + private taskType: string; - constructor({ googleApiKey, model, taskType }: { googleApiKey: string, model?: string, taskType?: string }) { - // we used to construct the client here, but we need to async import the types - // for the openai npm package, and the constructor can not be async - this.api_key = googleApiKey; - this.model = model || "embedding-001"; - this.taskType = taskType || "RETRIEVAL_DOCUMENT"; - } + constructor({ + googleApiKey, + model, + taskType, + }: { + googleApiKey: string; + model?: string; + taskType?: string; + }) { + // we used to construct the client here, but we need to async import the types + // for the openai npm package, and the constructor can not be async + this.api_key = googleApiKey; + this.model = model || "embedding-001"; + this.taskType = taskType || "RETRIEVAL_DOCUMENT"; + } - private async loadClient() { - if(this.googleGenAiApi) return; - try { - // eslint-disable-next-line global-require,import/no-extraneous-dependencies - const { googleGenAi } = await GoogleGenerativeAiEmbeddingFunction.import(); - googleGenAiApi = googleGenAi; - // googleGenAiApi.init(this.api_key); - googleGenAiApi = new googleGenAiApi(this.api_key); - } catch (_a) { - // @ts-ignore - if (_a.code === 'MODULE_NOT_FOUND') { - throw new Error("Please install the @google/generative-ai package to use the GoogleGenerativeAiEmbeddingFunction, `npm install -S @google/generative-ai`"); - } - throw _a; // Re-throw other errors - } - this.googleGenAiApi = googleGenAiApi; + private async loadClient() { + if (this.googleGenAiApi) return; + try { + // eslint-disable-next-line global-require,import/no-extraneous-dependencies + const { googleGenAi } = + await GoogleGenerativeAiEmbeddingFunction.import(); + googleGenAiApi = googleGenAi; + // googleGenAiApi.init(this.api_key); + googleGenAiApi = new googleGenAiApi(this.api_key); + } catch (_a) { + // @ts-ignore + if (_a.code === "MODULE_NOT_FOUND") { + throw new Error( + "Please install the @google/generative-ai package to use the GoogleGenerativeAiEmbeddingFunction, `npm install -S @google/generative-ai`", + ); + } + throw _a; // Re-throw other errors } + this.googleGenAiApi = googleGenAiApi; + } - public async generate(texts: string[]) { - - await this.loadClient(); - const model = this.googleGenAiApi.getGenerativeModel({ model: this.model}); - const response = await model.batchEmbedContents({ - requests: texts.map((t) => ({ - content: { parts: [{ text: t }] }, - taskType: this.taskType, - })), - }); - const embeddings = response.embeddings.map((e: any) => e.values); + public async generate(texts: string[]) { + await this.loadClient(); + const model = this.googleGenAiApi.getGenerativeModel({ model: this.model }); + const response = await model.batchEmbedContents({ + requests: texts.map((t) => ({ + content: { parts: [{ text: t }] }, + taskType: this.taskType, + })), + }); + const embeddings = response.embeddings.map((e: any) => e.values); - return embeddings; - } + return embeddings; + } - /** @ignore */ - static async import(): Promise<{ - // @ts-ignore - googleGenAi: typeof import("@google/generative-ai"); - }> { - try { - // @ts-ignore - const { GoogleGenerativeAI } = await import("@google/generative-ai"); - const googleGenAi = GoogleGenerativeAI; - return { googleGenAi }; - } catch (e) { - throw new Error( - "Please install @google/generative-ai as a dependency with, e.g. `yarn add @google/generative-ai`" - ); - } + /** @ignore */ + static async import(): Promise<{ + // @ts-ignore + googleGenAi: typeof import("@google/generative-ai"); + }> { + try { + // @ts-ignore + const { GoogleGenerativeAI } = await import("@google/generative-ai"); + const googleGenAi = GoogleGenerativeAI; + return { googleGenAi }; + } catch (e) { + throw new Error( + "Please install @google/generative-ai as a dependency with, e.g. `yarn add @google/generative-ai`", + ); } - + } } diff --git a/clients/js/src/embeddings/HuggingFaceEmbeddingServerFunction.ts b/clients/js/src/embeddings/HuggingFaceEmbeddingServerFunction.ts index dcbc62ecb70..b65c85f3f2f 100644 --- a/clients/js/src/embeddings/HuggingFaceEmbeddingServerFunction.ts +++ b/clients/js/src/embeddings/HuggingFaceEmbeddingServerFunction.ts @@ -3,29 +3,28 @@ import { IEmbeddingFunction } from "./IEmbeddingFunction"; let CohereAiApi: any; export class HuggingFaceEmbeddingServerFunction implements IEmbeddingFunction { - private url: string; + private url: string; - constructor({ url }: { url: string }) { - // we used to construct the client here, but we need to async import the types - // for the openai npm package, and the constructor can not be async - this.url = url; - } - - public async generate(texts: string[]) { - const response = await fetch(this.url, { - method: 'POST', - headers: { - 'Content-Type': 'application/json' - }, - body: JSON.stringify({ 'inputs': texts }) - }); + constructor({ url }: { url: string }) { + // we used to construct the client here, but we need to async import the types + // for the openai npm package, and the constructor can not be async + this.url = url; + } - if (!response.ok) { - throw new Error(`Failed to generate embeddings: ${response.statusText}`); - } + public async generate(texts: string[]) { + const response = await fetch(this.url, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ inputs: texts }), + }); - const data = await response.json(); - return data; + if (!response.ok) { + throw new Error(`Failed to generate embeddings: ${response.statusText}`); } + const data = await response.json(); + return data; + } } diff --git a/clients/js/src/embeddings/IEmbeddingFunction.ts b/clients/js/src/embeddings/IEmbeddingFunction.ts index dcc21ab19d9..7c3f85bc473 100644 --- a/clients/js/src/embeddings/IEmbeddingFunction.ts +++ b/clients/js/src/embeddings/IEmbeddingFunction.ts @@ -1,3 +1,3 @@ export interface IEmbeddingFunction { - generate(texts: string[]): Promise; + generate(texts: string[]): Promise; } diff --git a/clients/js/src/embeddings/JinaEmbeddingFunction.ts b/clients/js/src/embeddings/JinaEmbeddingFunction.ts index a91f94749f8..0ac92b2a9fa 100644 --- a/clients/js/src/embeddings/JinaEmbeddingFunction.ts +++ b/clients/js/src/embeddings/JinaEmbeddingFunction.ts @@ -5,20 +5,26 @@ export class JinaEmbeddingFunction implements IEmbeddingFunction { private api_url: string; private headers: { [key: string]: string }; - constructor({ jinaai_api_key, model_name }: { jinaai_api_key: string; model_name?: string }) { - this.model_name = model_name || 'jina-embeddings-v2-base-en'; - this.api_url = 'https://api.jina.ai/v1/embeddings'; + constructor({ + jinaai_api_key, + model_name, + }: { + jinaai_api_key: string; + model_name?: string; + }) { + this.model_name = model_name || "jina-embeddings-v2-base-en"; + this.api_url = "https://api.jina.ai/v1/embeddings"; this.headers = { Authorization: `Bearer ${jinaai_api_key}`, - 'Accept-Encoding': 'identity', - 'Content-Type': 'application/json', + "Accept-Encoding": "identity", + "Content-Type": "application/json", }; } public async generate(texts: string[]) { try { const response = await fetch(this.api_url, { - method: 'POST', + method: "POST", headers: this.headers, body: JSON.stringify({ input: texts, diff --git a/clients/js/src/embeddings/OllamaEmbeddingFunction.ts b/clients/js/src/embeddings/OllamaEmbeddingFunction.ts new file mode 100644 index 00000000000..bef8806f158 --- /dev/null +++ b/clients/js/src/embeddings/OllamaEmbeddingFunction.ts @@ -0,0 +1,34 @@ +import { IEmbeddingFunction } from "./IEmbeddingFunction"; + +export class OllamaEmbeddingFunction implements IEmbeddingFunction { + private readonly url: string; + private readonly model: string; + + constructor({ url, model }: { url: string, model: string }) { + // we used to construct the client here, but we need to async import the types + // for the openai npm package, and the constructor can not be async + this.url = url; + this.model = model; + } + + public async generate(texts: string[]) { + let embeddings:number[][] = []; + for (let text of texts) { + const response = await fetch(this.url, { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify({ 'model':this.model, 'prompt': text }) + }); + + if (!response.ok) { + throw new Error(`Failed to generate embeddings: ${response.status} (${response.statusText})`); + } + let finalResponse = await response.json(); + embeddings.push(finalResponse['embedding']); + } + return embeddings; + } + +} diff --git a/clients/js/src/embeddings/OpenAIEmbeddingFunction.ts b/clients/js/src/embeddings/OpenAIEmbeddingFunction.ts index 0b4be92eec1..bd61236b439 100644 --- a/clients/js/src/embeddings/OpenAIEmbeddingFunction.ts +++ b/clients/js/src/embeddings/OpenAIEmbeddingFunction.ts @@ -1,151 +1,157 @@ -import {IEmbeddingFunction} from "./IEmbeddingFunction"; +import { IEmbeddingFunction } from "./IEmbeddingFunction"; let OpenAIApi: any; let openAiVersion = null; let openAiMajorVersion = null; interface OpenAIAPI { - createEmbedding: (params: { - model: string; - input: string[]; - user?: string; - }) => Promise; + createEmbedding: (params: { + model: string; + input: string[]; + user?: string; + }) => Promise; } class OpenAIAPIv3 implements OpenAIAPI { - private readonly configuration: any; - private openai: any; - - constructor(configuration: { organization: string, apiKey: string }) { - this.configuration = new OpenAIApi.Configuration({ - organization: configuration.organization, - apiKey: configuration.apiKey, - }); - this.openai = new OpenAIApi.OpenAIApi(this.configuration); - } - - public async createEmbedding(params: { - model: string, - input: string[], - user?: string - }): Promise { - const embeddings: number[][] = []; - const response = await this.openai.createEmbedding({ - model: params.model, - input: params.input, - }).catch((error: any) => { - throw error; - }); - // @ts-ignore - const data = response.data["data"]; - for (let i = 0; i < data.length; i += 1) { - embeddings.push(data[i]["embedding"]); - } - return embeddings + private readonly configuration: any; + private openai: any; + + constructor(configuration: { organization: string; apiKey: string }) { + this.configuration = new OpenAIApi.Configuration({ + organization: configuration.organization, + apiKey: configuration.apiKey, + }); + this.openai = new OpenAIApi.OpenAIApi(this.configuration); + } + + public async createEmbedding(params: { + model: string; + input: string[]; + user?: string; + }): Promise { + const embeddings: number[][] = []; + const response = await this.openai + .createEmbedding({ + model: params.model, + input: params.input, + }) + .catch((error: any) => { + throw error; + }); + // @ts-ignore + const data = response.data["data"]; + for (let i = 0; i < data.length; i += 1) { + embeddings.push(data[i]["embedding"]); } + return embeddings; + } } class OpenAIAPIv4 implements OpenAIAPI { - private readonly apiKey: any; - private openai: any; - - constructor(apiKey: any) { - this.apiKey = apiKey; - this.openai = new OpenAIApi({ - apiKey: this.apiKey, - }); - } - - public async createEmbedding(params: { - model: string, - input: string[], - user?: string - }): Promise { - const embeddings: number[][] = []; - const response = await this.openai.embeddings.create(params); - const data = response["data"]; - for (let i = 0; i < data.length; i += 1) { - embeddings.push(data[i]["embedding"]); - } - return embeddings + private readonly apiKey: any; + private openai: any; + + constructor(apiKey: any) { + this.apiKey = apiKey; + this.openai = new OpenAIApi({ + apiKey: this.apiKey, + }); + } + + public async createEmbedding(params: { + model: string; + input: string[]; + user?: string; + }): Promise { + const embeddings: number[][] = []; + const response = await this.openai.embeddings.create(params); + const data = response["data"]; + for (let i = 0; i < data.length; i += 1) { + embeddings.push(data[i]["embedding"]); } + return embeddings; + } } export class OpenAIEmbeddingFunction implements IEmbeddingFunction { - private api_key: string; - private org_id: string; - private model: string; - private openaiApi?: OpenAIAPI; - - constructor({openai_api_key, openai_model, openai_organization_id}: { - openai_api_key: string, - openai_model?: string, - openai_organization_id?: string - }) { - // we used to construct the client here, but we need to async import the types - // for the openai npm package, and the constructor can not be async - this.api_key = openai_api_key; - this.org_id = openai_organization_id || ""; - this.model = openai_model || "text-embedding-ada-002"; - } - - private async loadClient() { - // cache the client - if(this.openaiApi) return; - - try { - const { openai, version } = await OpenAIEmbeddingFunction.import(); - OpenAIApi = openai; - let versionVar: string = version; - openAiVersion = versionVar.replace(/[^0-9.]/g, ''); - openAiMajorVersion = parseInt(openAiVersion.split('.')[0]); - } catch (_a) { - // @ts-ignore - if (_a.code === 'MODULE_NOT_FOUND') { - throw new Error("Please install the openai package to use the OpenAIEmbeddingFunction, `npm install -S openai`"); - } - throw _a; // Re-throw other errors - } - - if (openAiMajorVersion > 3) { - this.openaiApi = new OpenAIAPIv4(this.api_key); - } else { - this.openaiApi = new OpenAIAPIv3({ - organization: this.org_id, - apiKey: this.api_key, - }); - } + private api_key: string; + private org_id: string; + private model: string; + private openaiApi?: OpenAIAPI; + + constructor({ + openai_api_key, + openai_model, + openai_organization_id, + }: { + openai_api_key: string; + openai_model?: string; + openai_organization_id?: string; + }) { + // we used to construct the client here, but we need to async import the types + // for the openai npm package, and the constructor can not be async + this.api_key = openai_api_key; + this.org_id = openai_organization_id || ""; + this.model = openai_model || "text-embedding-ada-002"; + } + + private async loadClient() { + // cache the client + if (this.openaiApi) return; + + try { + const { openai, version } = await OpenAIEmbeddingFunction.import(); + OpenAIApi = openai; + let versionVar: string = version; + openAiVersion = versionVar.replace(/[^0-9.]/g, ""); + openAiMajorVersion = parseInt(openAiVersion.split(".")[0]); + } catch (_a) { + // @ts-ignore + if (_a.code === "MODULE_NOT_FOUND") { + throw new Error( + "Please install the openai package to use the OpenAIEmbeddingFunction, `npm install -S openai`", + ); + } + throw _a; // Re-throw other errors } - public async generate(texts: string[]): Promise { - - await this.loadClient(); - - return await this.openaiApi!.createEmbedding({ - model: this.model, - input: texts, - }).catch((error: any) => { - throw error; - }); + if (openAiMajorVersion > 3) { + this.openaiApi = new OpenAIAPIv4(this.api_key); + } else { + this.openaiApi = new OpenAIAPIv3({ + organization: this.org_id, + apiKey: this.api_key, + }); } - - /** @ignore */ - static async import(): Promise<{ - // @ts-ignore - openai: typeof import("openai"); - version: string; - }> { - try { - // @ts-ignore - const { default: openai } = await import("openai"); - // @ts-ignore - const { VERSION } = await import('openai/version'); - return { openai, version: VERSION }; - } catch (e) { - throw new Error( - "Please install openai as a dependency with, e.g. `yarn add openai`" - ); - } + } + + public async generate(texts: string[]): Promise { + await this.loadClient(); + + return await this.openaiApi!.createEmbedding({ + model: this.model, + input: texts, + }).catch((error: any) => { + throw error; + }); + } + + /** @ignore */ + static async import(): Promise<{ + // @ts-ignore + openai: typeof import("openai"); + version: string; + }> { + try { + // @ts-ignore + const { default: openai } = await import("openai"); + // @ts-ignore + const { VERSION } = await import("openai/version"); + return { openai, version: VERSION }; + } catch (e) { + throw new Error( + "Please install openai as a dependency with, e.g. `yarn add openai`", + ); } - + } } diff --git a/clients/js/src/embeddings/TransformersEmbeddingFunction.ts b/clients/js/src/embeddings/TransformersEmbeddingFunction.ts index aece174b03c..45a39d2a129 100644 --- a/clients/js/src/embeddings/TransformersEmbeddingFunction.ts +++ b/clients/js/src/embeddings/TransformersEmbeddingFunction.ts @@ -39,21 +39,21 @@ export class TransformersEmbeddingFunction implements IEmbeddingFunction { public async generate(texts: string[]): Promise { await this.loadClient(); - // Store a promise that resolves to the pipeline + // Store a promise that resolves to the pipeline this.pipelinePromise = new Promise(async (resolve, reject) => { try { - const pipeline = this.transformersApi + const pipeline = this.transformersApi; - const quantized = this.quantized - const revision = this.revision - const progress_callback = this.progress_callback + const quantized = this.quantized; + const revision = this.revision; + const progress_callback = this.progress_callback; resolve( await pipeline("feature-extraction", this.model, { quantized, revision, progress_callback, - }) + }), ); } catch (e) { reject(e); @@ -66,34 +66,36 @@ export class TransformersEmbeddingFunction implements IEmbeddingFunction { } private async loadClient() { - if(this.transformersApi) return; - try { - // eslint-disable-next-line global-require,import/no-extraneous-dependencies - let { pipeline } = await TransformersEmbeddingFunction.import(); - TransformersApi = pipeline; - } catch (_a) { - // @ts-ignore - if (_a.code === 'MODULE_NOT_FOUND') { - throw new Error("Please install the @xenova/transformers package to use the TransformersEmbeddingFunction, `npm install -S @xenova/transformers`"); - } - throw _a; // Re-throw other errors + if (this.transformersApi) return; + try { + // eslint-disable-next-line global-require,import/no-extraneous-dependencies + let { pipeline } = await TransformersEmbeddingFunction.import(); + TransformersApi = pipeline; + } catch (_a) { + // @ts-ignore + if (_a.code === "MODULE_NOT_FOUND") { + throw new Error( + "Please install the @xenova/transformers package to use the TransformersEmbeddingFunction, `npm install -S @xenova/transformers`", + ); } - this.transformersApi = TransformersApi; + throw _a; // Re-throw other errors + } + this.transformersApi = TransformersApi; } /** @ignore */ static async import(): Promise<{ - // @ts-ignore - pipeline: typeof import("@xenova/transformers"); + // @ts-ignore + pipeline: typeof import("@xenova/transformers"); }> { - try { - // @ts-ignore - const { pipeline } = await import("@xenova/transformers"); - return { pipeline }; - } catch (e) { - throw new Error( - "Please install @xenova/transformers as a dependency with, e.g. `yarn add @xenova/transformers`" - ); - } + try { + // @ts-ignore + const { pipeline } = await import("@xenova/transformers"); + return { pipeline }; + } catch (e) { + throw new Error( + "Please install @xenova/transformers as a dependency with, e.g. `yarn add @xenova/transformers`", + ); + } } } diff --git a/clients/js/src/index.ts b/clients/js/src/index.ts index 27316d1164a..c925f9e4871 100644 --- a/clients/js/src/index.ts +++ b/clients/js/src/index.ts @@ -1,45 +1,46 @@ -export { ChromaClient } from './ChromaClient'; -export { AdminClient } from './AdminClient'; -export { CloudClient } from './CloudClient'; -export { Collection } from './Collection'; +export { ChromaClient } from "./ChromaClient"; +export { AdminClient } from "./AdminClient"; +export { CloudClient } from "./CloudClient"; +export { Collection } from "./Collection"; +export { IEmbeddingFunction } from "./embeddings/IEmbeddingFunction"; +export { OpenAIEmbeddingFunction } from "./embeddings/OpenAIEmbeddingFunction"; +export { CohereEmbeddingFunction } from "./embeddings/CohereEmbeddingFunction"; +export { TransformersEmbeddingFunction } from "./embeddings/TransformersEmbeddingFunction"; +export { DefaultEmbeddingFunction } from "./embeddings/DefaultEmbeddingFunction"; +export { HuggingFaceEmbeddingServerFunction } from "./embeddings/HuggingFaceEmbeddingServerFunction"; +export { JinaEmbeddingFunction } from "./embeddings/JinaEmbeddingFunction"; +export { GoogleGenerativeAiEmbeddingFunction } from "./embeddings/GoogleGeminiEmbeddingFunction"; +export { OllamaEmbeddingFunction } from './embeddings/OllamaEmbeddingFunction'; -export { IEmbeddingFunction } from './embeddings/IEmbeddingFunction'; -export { OpenAIEmbeddingFunction } from './embeddings/OpenAIEmbeddingFunction'; -export { CohereEmbeddingFunction } from './embeddings/CohereEmbeddingFunction'; -export { TransformersEmbeddingFunction } from './embeddings/TransformersEmbeddingFunction'; -export { DefaultEmbeddingFunction } from './embeddings/DefaultEmbeddingFunction'; -export { HuggingFaceEmbeddingServerFunction } from './embeddings/HuggingFaceEmbeddingServerFunction'; -export { JinaEmbeddingFunction } from './embeddings/JinaEmbeddingFunction'; -export { GoogleGenerativeAiEmbeddingFunction } from './embeddings/GoogleGeminiEmbeddingFunction'; export { - IncludeEnum, - GetParams, - CollectionType, - CollectionMetadata, - Embedding, - Embeddings, - Metadata, - Metadatas, - Document, - Documents, - ID, - IDs, - Where, - WhereDocument, - GetResponse, - QueryResponse, - ListCollectionsParams, - ChromaClientParams, - CreateCollectionParams, - GetOrCreateCollectionParams, - GetCollectionParams, - DeleteCollectionParams, - AddParams, - UpsertParams, - UpdateParams, - ModifyCollectionParams, - QueryParams, - PeekParams, - DeleteParams -} from './types'; + IncludeEnum, + GetParams, + CollectionType, + CollectionMetadata, + Embedding, + Embeddings, + Metadata, + Metadatas, + Document, + Documents, + ID, + IDs, + Where, + WhereDocument, + GetResponse, + QueryResponse, + ListCollectionsParams, + ChromaClientParams, + CreateCollectionParams, + GetOrCreateCollectionParams, + GetCollectionParams, + DeleteCollectionParams, + AddParams, + UpsertParams, + UpdateParams, + ModifyCollectionParams, + QueryParams, + PeekParams, + DeleteParams, +} from "./types"; diff --git a/clients/js/src/types.ts b/clients/js/src/types.ts index 6c46d52c133..92e66f516c7 100644 --- a/clients/js/src/types.ts +++ b/clients/js/src/types.ts @@ -2,10 +2,10 @@ import { AuthOptions } from "./auth"; import { IEmbeddingFunction } from "./embeddings/IEmbeddingFunction"; export enum IncludeEnum { - Documents = 'documents', - Embeddings = 'embeddings', - Metadatas = 'metadatas', - Distances = 'distances' + Documents = "documents", + Embeddings = "embeddings", + Metadatas = "metadatas", + Distances = "distances", } type Number = number; @@ -31,7 +31,9 @@ type InclusionOperator = "$in" | "$nin"; type WhereOperator = "$gt" | "$gte" | "$lt" | "$lte" | "$ne" | "$eq"; type OperatorExpression = { - [key in WhereOperator | InclusionOperator | LogicalOperator ]?: LiteralValue | ListLiteralValue; + [key in WhereOperator | InclusionOperator | LogicalOperator]?: + | LiteralValue + | ListLiteralValue; }; type BaseWhere = { @@ -44,10 +46,13 @@ type LogicalWhere = { export type Where = BaseWhere | LogicalWhere; -type WhereDocumentOperator = "$contains" | LogicalOperator; +type WhereDocumentOperator = "$contains" | "$not_contains" | LogicalOperator; export type WhereDocument = { - [key in WhereDocumentOperator]?: LiteralValue | LiteralNumber | WhereDocument[]; + [key in WhereDocumentOperator]?: + | LiteralValue + | LiteralNumber + | WhereDocument[]; }; export type CollectionType = { @@ -70,11 +75,11 @@ export type QueryResponse = { documents: (null | Document)[][]; metadatas: (null | Metadata)[][]; distances: null | number[][]; -} +}; export type AddResponse = { error: string; -} +}; export type CollectionMetadata = Record; @@ -85,72 +90,72 @@ export type ConfigOptions = { }; export type GetParams = { - ids?: ID | IDs, - where?: Where, - limit?: PositiveInteger, - offset?: PositiveInteger, - include?: IncludeEnum[], - whereDocument?: WhereDocument -} + ids?: ID | IDs; + where?: Where; + limit?: PositiveInteger; + offset?: PositiveInteger; + include?: IncludeEnum[]; + whereDocument?: WhereDocument; +}; export type ListCollectionsParams = { - limit?: PositiveInteger, - offset?: PositiveInteger, -} + limit?: PositiveInteger; + offset?: PositiveInteger; +}; export type ChromaClientParams = { - path?: string, - fetchOptions?: RequestInit, - auth?: AuthOptions, - tenant?: string, - database?: string, -} + path?: string; + fetchOptions?: RequestInit; + auth?: AuthOptions; + tenant?: string; + database?: string; +}; export type CreateCollectionParams = { - name: string, - metadata?: CollectionMetadata, - embeddingFunction?: IEmbeddingFunction -} + name: string; + metadata?: CollectionMetadata; + embeddingFunction?: IEmbeddingFunction; +}; -export type GetOrCreateCollectionParams = CreateCollectionParams +export type GetOrCreateCollectionParams = CreateCollectionParams; export type GetCollectionParams = { name: string; - embeddingFunction?: IEmbeddingFunction -} + embeddingFunction?: IEmbeddingFunction; +}; export type DeleteCollectionParams = { - name: string -} + name: string; +}; export type AddParams = { - ids: ID | IDs, - embeddings?: Embedding | Embeddings, - metadatas?: Metadata | Metadatas, - documents?: Document | Documents, -} + ids: ID | IDs; + embeddings?: Embedding | Embeddings; + metadatas?: Metadata | Metadatas; + documents?: Document | Documents; +}; export type UpsertParams = AddParams; export type UpdateParams = AddParams; export type ModifyCollectionParams = { - name?: string, - metadata?: CollectionMetadata -} + name?: string; + metadata?: CollectionMetadata; +}; export type QueryParams = { - queryEmbeddings?: Embedding | Embeddings, - nResults?: PositiveInteger, - where?: Where, - queryTexts?: string | string[], - whereDocument?: WhereDocument, // {"$contains":"search_string"} - include?: IncludeEnum[] // ["metadata", "document"] -} + queryEmbeddings?: Embedding | Embeddings; + nResults?: PositiveInteger; + where?: Where; + queryTexts?: string | string[]; + whereDocument?: WhereDocument; // {"$contains":"search_string"} + include?: IncludeEnum[]; // ["metadata", "document"] +}; -export type PeekParams = { limit?: PositiveInteger } +export type PeekParams = { limit?: PositiveInteger }; export type DeleteParams = { - ids?: ID | IDs, - where?: Where, - whereDocument?: WhereDocument -} + ids?: ID | IDs; + where?: Where; + whereDocument?: WhereDocument; +}; diff --git a/clients/js/src/utils.ts b/clients/js/src/utils.ts index e3ad5361e61..cf43c18ccb3 100644 --- a/clients/js/src/utils.ts +++ b/clients/js/src/utils.ts @@ -13,7 +13,7 @@ export function toArray(obj: T | Array): Array { // a function to convert an array to array of arrays export function toArrayOfArrays( - obj: Array> | Array + obj: Array> | Array, ): Array> { if (Array.isArray(obj[0])) { return obj as Array>; @@ -56,7 +56,7 @@ export async function handleError(error: unknown) { } export async function handleSuccess( - response: Response | string | Count200Response + response: Response | string | Count200Response, ) { switch (true) { case response instanceof Response: @@ -83,17 +83,24 @@ export async function importOptionalModule(moduleName: string) { return Function(`return import("${moduleName}")`)(); } +export async function validateTenantDatabase( + adminClient: AdminClient, + tenant: string, + database: string, +): Promise { + try { + await adminClient.getTenant({ name: tenant }); + } catch (error) { + throw new Error( + `Error: ${error}, Could not connect to tenant ${tenant}. Are you sure it exists?`, + ); + } -export async function validateTenantDatabase(adminClient: AdminClient, tenant: string, database: string): Promise { - try { - await adminClient.getTenant({name: tenant}); - } catch (error) { - throw new Error(`Error: ${error}, Could not connect to tenant ${tenant}. Are you sure it exists?`); - } - - try { - await adminClient.getDatabase({name: database, tenantName: tenant}); - } catch (error) { - throw new Error(`Error: ${error}, Could not connect to database ${database} for tenant ${tenant}. Are you sure it exists?`); - } + try { + await adminClient.getDatabase({ name: database, tenantName: tenant }); + } catch (error) { + throw new Error( + `Error: ${error}, Could not connect to database ${database} for tenant ${tenant}. Are you sure it exists?`, + ); + } } diff --git a/clients/js/test/add.collections.test.ts b/clients/js/test/add.collections.test.ts index cb89fa8dbe0..41b3de3fef5 100644 --- a/clients/js/test/add.collections.test.ts +++ b/clients/js/test/add.collections.test.ts @@ -1,10 +1,11 @@ -import { expect, test } from '@jest/globals'; -import chroma from './initClient' -import { DOCUMENTS, EMBEDDINGS, IDS } from './data'; -import { METADATAS } from './data'; +import { expect, test } from "@jest/globals"; +import chroma from "./initClient"; +import { DOCUMENTS, EMBEDDINGS, IDS } from "./data"; +import { METADATAS } from "./data"; import { IncludeEnum } from "../src/types"; -import {OpenAIEmbeddingFunction} from "../src/embeddings/OpenAIEmbeddingFunction"; -import {CohereEmbeddingFunction} from "../src/embeddings/CohereEmbeddingFunction"; +import { OpenAIEmbeddingFunction } from "../src/embeddings/OpenAIEmbeddingFunction"; +import { CohereEmbeddingFunction } from "../src/embeddings/CohereEmbeddingFunction"; +import { OllamaEmbeddingFunction } from "../src/embeddings/OllamaEmbeddingFunction"; test("it should add single embeddings to a collection", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); @@ -15,9 +16,8 @@ test("it should add single embeddings to a collection", async () => { const count = await collection.count(); expect(count).toBe(1); var res = await collection.get({ - ids: [ids], include: [ - IncludeEnum.Embeddings, - ] + ids: [ids], + include: [IncludeEnum.Embeddings], }); expect(res.embeddings![0]).toEqual(embeddings); }); @@ -29,51 +29,55 @@ test("it should add batch embeddings to a collection", async () => { const count = await collection.count(); expect(count).toBe(3); var res = await collection.get({ - ids: IDS, include: [ - IncludeEnum.Embeddings, - ] + ids: IDS, + include: [IncludeEnum.Embeddings], }); expect(res.embeddings).toEqual(EMBEDDINGS); // reverse because of the order of the ids }); - if (!process.env.OPENAI_API_KEY) { - test.skip("it should add OpenAI embeddings", async () => { - }); + test.skip("it should add OpenAI embeddings", async () => {}); } else { test("it should add OpenAI embeddings", async () => { await chroma.reset(); - const embedder = new OpenAIEmbeddingFunction({ openai_api_key: process.env.OPENAI_API_KEY || "" }) - const collection = await chroma.createCollection({ name: "test" ,embeddingFunction: embedder}); + const embedder = new OpenAIEmbeddingFunction({ + openai_api_key: process.env.OPENAI_API_KEY || "", + }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embedder, + }); const embeddings = await embedder.generate(DOCUMENTS); await collection.add({ ids: IDS, embeddings: embeddings }); const count = await collection.count(); expect(count).toBe(3); var res = await collection.get({ - ids: IDS, include: [ - IncludeEnum.Embeddings, - ] + ids: IDS, + include: [IncludeEnum.Embeddings], }); expect(res.embeddings).toEqual(embeddings); // reverse because of the order of the ids }); } if (!process.env.COHERE_API_KEY) { - test.skip("it should add Cohere embeddings", async () => { - }); + test.skip("it should add Cohere embeddings", async () => {}); } else { test("it should add Cohere embeddings", async () => { await chroma.reset(); - const embedder = new CohereEmbeddingFunction({ cohere_api_key: process.env.COHERE_API_KEY || "" }) - const collection = await chroma.createCollection({ name: "test" ,embeddingFunction: embedder}); + const embedder = new CohereEmbeddingFunction({ + cohere_api_key: process.env.COHERE_API_KEY || "", + }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embedder, + }); const embeddings = await embedder.generate(DOCUMENTS); await collection.add({ ids: IDS, embeddings: embeddings }); const count = await collection.count(); expect(count).toBe(3); var res = await collection.get({ - ids: IDS, include: [ - IncludeEnum.Embeddings, - ] + ids: IDS, + include: [IncludeEnum.Embeddings], }); expect(res.embeddings).toEqual(embeddings); // reverse because of the order of the ids }); @@ -82,21 +86,65 @@ if (!process.env.COHERE_API_KEY) { test("add documents", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - let resp = await collection.add({ ids: IDS, embeddings: EMBEDDINGS, documents: DOCUMENTS }); - expect(resp).toBe(true) + let resp = await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + documents: DOCUMENTS, + }); + expect(resp).toBe(true); const results = await collection.get({ ids: ["test1"] }); expect(results.documents[0]).toBe("This is a test"); }); -test('It should return an error when inserting duplicate IDs in the same batch', async () => { - await chroma.reset() +test("It should return an error when inserting duplicate IDs in the same batch", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = IDS.concat(["test1"]); + const embeddings = EMBEDDINGS.concat([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]); + const metadatas = METADATAS.concat([{ test: "test1", float_value: 0.1 }]); + try { + await collection.add({ ids, embeddings, metadatas }); + } catch (e: any) { + expect(e.message).toMatch("duplicates"); + } +}); + +test("should error on empty embedding", async () => { + await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - const ids = IDS.concat(["test1"]) - const embeddings = EMBEDDINGS.concat([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]) - const metadatas = METADATAS.concat([{ test: 'test1', 'float_value': 0.1 }]) + const ids = ["id1"]; + const embeddings = [[]]; + const metadatas = [{ test: "test1", float_value: 0.1 }]; try { await collection.add({ ids, embeddings, metadatas }); } catch (e: any) { - expect(e.message).toMatch('duplicates') + expect(e.message).toMatch("got empty embedding at pos"); } -}) +}); + +if (!process.env.OLLAMA_SERVER_URL) { + test.skip("it should use ollama EF, OLLAMA_SERVER_URL not defined", async () => {}); +} else { + test("it should use ollama EF", async () => { + await chroma.reset(); + const embedder = new OllamaEmbeddingFunction({ + url: + process.env.OLLAMA_SERVER_URL || + "http://127.0.0.1:11434/api/embeddings", + model: "nomic-embed-text", + }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embedder, + }); + const embeddings = await embedder.generate(DOCUMENTS); + await collection.add({ ids: IDS, embeddings: embeddings }); + const count = await collection.count(); + expect(count).toBe(3); + var res = await collection.get({ + ids: IDS, + include: [IncludeEnum.Embeddings], + }); + expect(res.embeddings).toEqual(embeddings); // reverse because of the order of the ids + }); +} diff --git a/clients/js/test/admin.test.ts b/clients/js/test/admin.test.ts index d0ee72db8c4..a0ba58b7fd9 100644 --- a/clients/js/test/admin.test.ts +++ b/clients/js/test/admin.test.ts @@ -3,48 +3,64 @@ import { AdminClient } from "../src/AdminClient"; import adminClient from "./initAdminClient"; test("it should create the admin client connection", async () => { - expect(adminClient).toBeDefined(); - expect(adminClient).toBeInstanceOf(AdminClient); + expect(adminClient).toBeDefined(); + expect(adminClient).toBeInstanceOf(AdminClient); }); test("it should create and get a tenant", async () => { - await adminClient.createTenant({ name: "testTenant" }); - const tenant = await adminClient.getTenant({ name: "testTenant" }); - expect(tenant).toBeDefined(); - expect(tenant).toHaveProperty('name') - expect(tenant.name).toBe("testTenant"); -}) + await adminClient.createTenant({ name: "testTenant" }); + const tenant = await adminClient.getTenant({ name: "testTenant" }); + expect(tenant).toBeDefined(); + expect(tenant).toHaveProperty("name"); + expect(tenant.name).toBe("testTenant"); +}); test("it should create and get a database for a tenant", async () => { - await adminClient.createTenant({ name: "test3" }); - const database = await adminClient.createDatabase({ name: "test", tenantName: "test3" }); - expect(database).toBeDefined(); - expect(database).toHaveProperty('name') - expect(database.name).toBe("test"); - - const getDatabase = await adminClient.getDatabase({ name: "test", tenantName: "test3" }); - expect(getDatabase).toBeDefined(); - expect(getDatabase).toHaveProperty('name') - expect(getDatabase.name).toBe("test"); -}) + await adminClient.createTenant({ name: "test3" }); + const database = await adminClient.createDatabase({ + name: "test", + tenantName: "test3", + }); + expect(database).toBeDefined(); + expect(database).toHaveProperty("name"); + expect(database.name).toBe("test"); + + const getDatabase = await adminClient.getDatabase({ + name: "test", + tenantName: "test3", + }); + expect(getDatabase).toBeDefined(); + expect(getDatabase).toHaveProperty("name"); + expect(getDatabase.name).toBe("test"); +}); // test that it can set the tenant and database test("it should set the tenant and database", async () => { - // doesnt exist so should throw - await expect(adminClient.setTenant({ tenant: "testTenant", database: "testDatabase" })).rejects.toThrow(); + // doesnt exist so should throw + await expect( + adminClient.setTenant({ tenant: "testTenant", database: "testDatabase" }), + ).rejects.toThrow(); - await adminClient.createTenant({ name: "testTenant!" }); - await adminClient.createDatabase({ name: "test3!", tenantName: "testTenant!" }); + await adminClient.createTenant({ name: "testTenant!" }); + await adminClient.createDatabase({ + name: "test3!", + tenantName: "testTenant!", + }); - await adminClient.setTenant({ tenant: "testTenant!", database: "test3!" }); - expect(adminClient.tenant).toBe("testTenant!"); - expect(adminClient.database).toBe("test3!"); + await adminClient.setTenant({ tenant: "testTenant!", database: "test3!" }); + expect(adminClient.tenant).toBe("testTenant!"); + expect(adminClient.database).toBe("test3!"); - // doesnt exist so should throw - await expect(adminClient.setDatabase({database: "testDatabase2"})).rejects.toThrow(); + // doesnt exist so should throw + await expect( + adminClient.setDatabase({ database: "testDatabase2" }), + ).rejects.toThrow(); - await adminClient.createDatabase({ name: "testDatabase2", tenantName: "testTenant!" }); - await adminClient.setDatabase({database: "testDatabase2"}) + await adminClient.createDatabase({ + name: "testDatabase2", + tenantName: "testTenant!", + }); + await adminClient.setDatabase({ database: "testDatabase2" }); - expect(adminClient.database).toBe("testDatabase2"); -}) + expect(adminClient.database).toBe("testDatabase2"); +}); diff --git a/clients/js/test/auth.basic.test.ts b/clients/js/test/auth.basic.test.ts index 6253bb758a3..16e3b084126 100644 --- a/clients/js/test/auth.basic.test.ts +++ b/clients/js/test/auth.basic.test.ts @@ -1,33 +1,33 @@ -import {expect, test} from "@jest/globals"; -import {chromaBasic} from "./initClientWithAuth"; +import { expect, test } from "@jest/globals"; +import { chromaBasic } from "./initClientWithAuth"; import chromaNoAuth from "./initClient"; import { ChromaClient } from "../src/ChromaClient"; test("it should get the version without auth needed", async () => { - const version = await chromaNoAuth.version(); - expect(version).toBeDefined(); - expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/); + const version = await chromaNoAuth.version(); + expect(version).toBeDefined(); + expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/); }); test("it should get the heartbeat without auth needed", async () => { - const heartbeat = await chromaNoAuth.heartbeat(); - expect(heartbeat).toBeDefined(); - expect(heartbeat).toBeGreaterThan(0); + const heartbeat = await chromaNoAuth.heartbeat(); + expect(heartbeat).toBeDefined(); + expect(heartbeat).toBeGreaterThan(0); }); test("it should raise error when non authenticated", async () => { - await expect(chromaNoAuth.listCollections()).rejects.toMatchObject({ - status: 401 - }); + await expect(chromaNoAuth.listCollections()).rejects.toMatchObject({ + status: 401, + }); }); -test('it should list collections', async () => { - await chromaBasic.reset() - let collections = await chromaBasic.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - await chromaBasic.createCollection({name: "test"}); - collections = await chromaBasic.listCollections() - expect(collections.length).toBe(1) -}) +test("it should list collections", async () => { + await chromaBasic.reset(); + let collections = await chromaBasic.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + await chromaBasic.createCollection({ name: "test" }); + collections = await chromaBasic.listCollections(); + expect(collections.length).toBe(1); +}); diff --git a/clients/js/test/auth.token.test.ts b/clients/js/test/auth.token.test.ts index a57ea09e1f5..3d371faed0b 100644 --- a/clients/js/test/auth.token.test.ts +++ b/clients/js/test/auth.token.test.ts @@ -1,70 +1,79 @@ -import {expect, test} from "@jest/globals"; -import {ChromaClient} from "../src/ChromaClient"; -import {chromaTokenDefault, chromaTokenBearer, chromaTokenXToken, cloudClient} from "./initClientWithAuth"; +import { expect, test } from "@jest/globals"; +import { ChromaClient } from "../src/ChromaClient"; +import { + chromaTokenDefault, + chromaTokenBearer, + chromaTokenXToken, + cloudClient, +} from "./initClientWithAuth"; import chromaNoAuth from "./initClient"; test("it should get the version without auth needed", async () => { - const version = await chromaNoAuth.version(); - expect(version).toBeDefined(); - expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/); + const version = await chromaNoAuth.version(); + expect(version).toBeDefined(); + expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/); }); test("it should get the heartbeat without auth needed", async () => { - const heartbeat = await chromaNoAuth.heartbeat(); - expect(heartbeat).toBeDefined(); - expect(heartbeat).toBeGreaterThan(0); + const heartbeat = await chromaNoAuth.heartbeat(); + expect(heartbeat).toBeDefined(); + expect(heartbeat).toBeGreaterThan(0); }); test("it should raise error when non authenticated", async () => { - await expect(chromaNoAuth.listCollections()).rejects.toMatchObject({ - status: 401 - }); + await expect(chromaNoAuth.listCollections()).rejects.toMatchObject({ + status: 401, + }); }); if (!process.env.XTOKEN_TEST) { - test('it should list collections with default token config', async () => { - await chromaTokenDefault.reset() - let collections = await chromaTokenDefault.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - const collection = await chromaTokenDefault.createCollection({name: "test"}); - collections = await chromaTokenDefault.listCollections() - expect(collections.length).toBe(1) - }) + test("it should list collections with default token config", async () => { + await chromaTokenDefault.reset(); + let collections = await chromaTokenDefault.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + const collection = await chromaTokenDefault.createCollection({ + name: "test", + }); + collections = await chromaTokenDefault.listCollections(); + expect(collections.length).toBe(1); + }); - test('it should list collections with explicit bearer token config', async () => { - await chromaTokenBearer.reset() - let collections = await chromaTokenBearer.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - const collection = await chromaTokenBearer.createCollection({name: "test"}); - collections = await chromaTokenBearer.listCollections() - expect(collections.length).toBe(1) - }) + test("it should list collections with explicit bearer token config", async () => { + await chromaTokenBearer.reset(); + let collections = await chromaTokenBearer.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + const collection = await chromaTokenBearer.createCollection({ + name: "test", + }); + collections = await chromaTokenBearer.listCollections(); + expect(collections.length).toBe(1); + }); } else { + test("it should list collections with explicit x-token token config", async () => { + await chromaTokenXToken.reset(); + let collections = await chromaTokenXToken.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + const collection = await chromaTokenXToken.createCollection({ + name: "test", + }); + collections = await chromaTokenXToken.listCollections(); + expect(collections.length).toBe(1); + }); - test('it should list collections with explicit x-token token config', async () => { - await chromaTokenXToken.reset() - let collections = await chromaTokenXToken.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - const collection = await chromaTokenXToken.createCollection({name: "test"}); - collections = await chromaTokenXToken.listCollections() - expect(collections.length).toBe(1) - }) - - test('it should list collections with explicit x-token token config in CloudClient', async () => { - await cloudClient.reset() - let collections = await cloudClient.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - const collection = await cloudClient.createCollection({name: "test"}); - collections = await cloudClient.listCollections() - expect(collections.length).toBe(1) - }) - + test("it should list collections with explicit x-token token config in CloudClient", async () => { + await cloudClient.reset(); + let collections = await cloudClient.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + const collection = await cloudClient.createCollection({ name: "test" }); + collections = await cloudClient.listCollections(); + expect(collections.length).toBe(1); + }); } diff --git a/clients/js/test/client.test.ts b/clients/js/test/client.test.ts index 512237a2457..d9d85486d0d 100644 --- a/clients/js/test/client.test.ts +++ b/clients/js/test/client.test.ts @@ -3,193 +3,198 @@ import { ChromaClient } from "../src/ChromaClient"; import chroma from "./initClient"; test("it should create the client connection", async () => { - expect(chroma).toBeDefined(); - expect(chroma).toBeInstanceOf(ChromaClient); + expect(chroma).toBeDefined(); + expect(chroma).toBeInstanceOf(ChromaClient); }); test("it should get the version", async () => { - const version = await chroma.version(); - expect(version).toBeDefined(); - expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/); + const version = await chroma.version(); + expect(version).toBeDefined(); + expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/); }); test("it should get the heartbeat", async () => { - const heartbeat = await chroma.heartbeat(); - expect(heartbeat).toBeDefined(); - expect(heartbeat).toBeGreaterThan(0); + const heartbeat = await chroma.heartbeat(); + expect(heartbeat).toBeDefined(); + expect(heartbeat).toBeGreaterThan(0); }); test("it should reset the database", async () => { - await chroma.reset(); - const collections = await chroma.listCollections(); - expect(collections).toBeDefined(); - expect(collections).toBeInstanceOf(Array); - expect(collections.length).toBe(0); - - const collection = await chroma.createCollection({ name: "test" }); - const collections2 = await chroma.listCollections(); - expect(collections2).toBeDefined(); - expect(collections2).toBeInstanceOf(Array); - expect(collections2.length).toBe(1); - - await chroma.reset(); - const collections3 = await chroma.listCollections(); - expect(collections3).toBeDefined(); - expect(collections3).toBeInstanceOf(Array); - expect(collections3.length).toBe(0); + await chroma.reset(); + const collections = await chroma.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + + const collection = await chroma.createCollection({ name: "test" }); + const collections2 = await chroma.listCollections(); + expect(collections2).toBeDefined(); + expect(collections2).toBeInstanceOf(Array); + expect(collections2.length).toBe(1); + + await chroma.reset(); + const collections3 = await chroma.listCollections(); + expect(collections3).toBeDefined(); + expect(collections3).toBeInstanceOf(Array); + expect(collections3.length).toBe(0); }); -test('it should list collections', async () => { - await chroma.reset() - let collections = await chroma.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - const collection = await chroma.createCollection({ name: "test" }); - collections = await chroma.listCollections() - expect(collections.length).toBe(1) -}) - -test('it should get a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const collection2 = await chroma.getCollection({ name: "test" }); - expect(collection).toBeDefined() - expect(collection2).toBeDefined() - expect(collection).toHaveProperty('name') - expect(collection2).toHaveProperty('name') - expect(collection.name).toBe(collection2.name) - expect(collection).toHaveProperty('id') - expect(collection2).toHaveProperty('id') - expect(collection.id).toBe(collection2.id) -}) - -test('it should delete a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - let collections = await chroma.listCollections() - expect(collections.length).toBe(1) - var resp = await chroma.deleteCollection({ name: "test" }); - collections = await chroma.listCollections() - expect(collections.length).toBe(0) -}) - -test('it should add single embeddings to a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = 'test1' - const embeddings = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - const metadatas = { test: 'test' } - await collection.add({ ids, embeddings, metadatas }) - const count = await collection.count() - expect(count).toBe(1) -}) - -test('it should add batch embeddings to a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add({ ids, embeddings }) - const count = await collection.count() - expect(count).toBe(3) -}) - -test('it should query a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add({ ids, embeddings }) - const results = await collection.query({ queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], nResults: 2 }) - expect(results).toBeDefined() - expect(results).toBeInstanceOf(Object) - // expect(results.embeddings[0].length).toBe(2) - const result: string[] = ['test1', 'test2'] - expect(result).toEqual(expect.arrayContaining(results.ids[0])); - expect(['test3']).not.toEqual(expect.arrayContaining(results.ids[0])); -}) - -test('it should peek a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add({ ids, embeddings }) - const results = await collection.peek({ limit: 2 }) - expect(results).toBeDefined() - expect(results).toBeInstanceOf(Object) - expect(results.ids.length).toBe(2) - expect(['test1', 'test2']).toEqual(expect.arrayContaining(results.ids)); -}) - -test('it should get a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - const metadatas = [{ test: 'test1' }, { test: 'test2' }, { test: 'test3' }] - await collection.add({ ids, embeddings, metadatas }) - const results = await collection.get({ ids: ['test1'] }) - expect(results).toBeDefined() - expect(results).toBeInstanceOf(Object) - expect(results.ids.length).toBe(1) - expect(['test1']).toEqual(expect.arrayContaining(results.ids)); - expect(['test2']).not.toEqual(expect.arrayContaining(results.ids)); - - const results2 = await collection.get({ where: { 'test': 'test1' } }) - expect(results2).toBeDefined() - expect(results2).toBeInstanceOf(Object) - expect(results2.ids.length).toBe(1) -}) - -test('it should delete a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - const metadatas = [{ test: 'test1' }, { test: 'test2' }, { test: 'test3' }] - await collection.add({ ids, embeddings, metadatas }) - let count = await collection.count() - expect(count).toBe(3) - var resp = await collection.delete({ where: { 'test': 'test1' } }) - count = await collection.count() - expect(count).toBe(2) -}) - -test('wrong code returns an error', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - const metadatas = [{ test: 'test1' }, { test: 'test2' }, { test: 'test3' }] - await collection.add({ ids, embeddings, metadatas }) +test("it should list collections", async () => { + await chroma.reset(); + let collections = await chroma.listCollections(); + expect(collections).toBeDefined(); + expect(collections).toBeInstanceOf(Array); + expect(collections.length).toBe(0); + const collection = await chroma.createCollection({ name: "test" }); + collections = await chroma.listCollections(); + expect(collections.length).toBe(1); +}); + +test("it should get a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const collection2 = await chroma.getCollection({ name: "test" }); + expect(collection).toBeDefined(); + expect(collection2).toBeDefined(); + expect(collection).toHaveProperty("name"); + expect(collection2).toHaveProperty("name"); + expect(collection.name).toBe(collection2.name); + expect(collection).toHaveProperty("id"); + expect(collection2).toHaveProperty("id"); + expect(collection.id).toBe(collection2.id); +}); + +test("it should delete a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + let collections = await chroma.listCollections(); + expect(collections.length).toBe(1); + var resp = await chroma.deleteCollection({ name: "test" }); + collections = await chroma.listCollections(); + expect(collections.length).toBe(0); +}); + +test("it should add single embeddings to a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = "test1"; + const embeddings = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + const metadatas = { test: "test" }; + await collection.add({ ids, embeddings, metadatas }); + const count = await collection.count(); + expect(count).toBe(1); +}); + +test("it should add batch embeddings to a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2", "test3"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + await collection.add({ ids, embeddings }); + const count = await collection.count(); + expect(count).toBe(3); +}); + +test("it should query a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2", "test3"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + await collection.add({ ids, embeddings }); + const results = await collection.query({ + queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + nResults: 2, + }); + expect(results).toBeDefined(); + expect(results).toBeInstanceOf(Object); + // expect(results.embeddings[0].length).toBe(2) + const result: string[] = ["test1", "test2"]; + expect(result).toEqual(expect.arrayContaining(results.ids[0])); + expect(["test3"]).not.toEqual(expect.arrayContaining(results.ids[0])); +}); + +test("it should peek a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2", "test3"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + await collection.add({ ids, embeddings }); + const results = await collection.peek({ limit: 2 }); + expect(results).toBeDefined(); + expect(results).toBeInstanceOf(Object); + expect(results.ids.length).toBe(2); + expect(["test1", "test2"]).toEqual(expect.arrayContaining(results.ids)); +}); + +test("it should get a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2", "test3"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + const metadatas = [{ test: "test1" }, { test: "test2" }, { test: "test3" }]; + await collection.add({ ids, embeddings, metadatas }); + const results = await collection.get({ ids: ["test1"] }); + expect(results).toBeDefined(); + expect(results).toBeInstanceOf(Object); + expect(results.ids.length).toBe(1); + expect(["test1"]).toEqual(expect.arrayContaining(results.ids)); + expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids)); + + const results2 = await collection.get({ where: { test: "test1" } }); + expect(results2).toBeDefined(); + expect(results2).toBeInstanceOf(Object); + expect(results2.ids.length).toBe(1); +}); + +test("it should delete a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2", "test3"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + const metadatas = [{ test: "test1" }, { test: "test2" }, { test: "test3" }]; + await collection.add({ ids, embeddings, metadatas }); + let count = await collection.count(); + expect(count).toBe(3); + var resp = await collection.delete({ where: { test: "test1" } }); + count = await collection.count(); + expect(count).toBe(2); +}); + +test("wrong code returns an error", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2", "test3"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + const metadatas = [{ test: "test1" }, { test: "test2" }, { test: "test3" }]; + await collection.add({ ids, embeddings, metadatas }); + const results = await collection.get({ // @ts-ignore - supposed to fail - const results = await collection.get({ where: { "test": { "$contains": "hello" } } }); - expect(results.error).toBeDefined() - expect(results.error).toContain("ValueError('Expected where operator") -}) + where: { test: { $contains: "hello" } }, + }); + expect(results.error).toBeDefined(); + expect(results.error).toContain("ValueError('Expected where operator"); +}); diff --git a/clients/js/test/collection.client.test.ts b/clients/js/test/collection.client.test.ts index 0045067c7ea..660808504dc 100644 --- a/clients/js/test/collection.client.test.ts +++ b/clients/js/test/collection.client.test.ts @@ -19,29 +19,44 @@ test("it should create a collection", async () => { const collection = await chroma.createCollection({ name: "test" }); expect(collection).toBeDefined(); expect(collection).toHaveProperty("name"); - expect(collection).toHaveProperty('id') + expect(collection).toHaveProperty("id"); expect(collection.name).toBe("test"); let collections = await chroma.listCollections(); - expect([{ name: "test", metadata: null, id: collection.id, database: "default_database", tenant: "default_tenant" }]).toEqual( - expect.arrayContaining(collections) - ); + expect([ + { + name: "test", + metadata: null, + id: collection.id, + database: "default_database", + tenant: "default_tenant", + }, + ]).toEqual(expect.arrayContaining(collections)); expect([{ name: "test2", metadata: null }]).not.toEqual( - expect.arrayContaining(collections) + expect.arrayContaining(collections), ); await chroma.reset(); - const collection2 = await chroma.createCollection({ name: "test2", metadata: { test: "test" } }); + const collection2 = await chroma.createCollection({ + name: "test2", + metadata: { test: "test" }, + }); expect(collection2).toBeDefined(); expect(collection2).toHaveProperty("name"); - expect(collection2).toHaveProperty('id') + expect(collection2).toHaveProperty("id"); expect(collection2.name).toBe("test2"); expect(collection2).toHaveProperty("metadata"); expect(collection2.metadata).toHaveProperty("test"); expect(collection2.metadata).toEqual({ test: "test" }); let collections2 = await chroma.listCollections(); - expect([{ name: "test2", metadata: { test: "test" }, id: collection2.id, database: "default_database", tenant: "default_tenant" }]).toEqual( - expect.arrayContaining(collections2) - ); + expect([ + { + name: "test2", + metadata: { test: "test" }, + id: collection2.id, + database: "default_database", + tenant: "default_tenant", + }, + ]).toEqual(expect.arrayContaining(collections2)); }); test("it should get a collection", async () => { diff --git a/clients/js/test/collection.test.ts b/clients/js/test/collection.test.ts index 4e4919d4932..b6604c48356 100644 --- a/clients/js/test/collection.test.ts +++ b/clients/js/test/collection.test.ts @@ -24,7 +24,7 @@ test("it should modify collection", async () => { const collection3 = await chroma.createCollection({ name: original_name, - metadata: original_metadata + metadata: original_metadata, }); expect(collection3.name).toBe(original_name); expect(collection3.metadata).toEqual(original_metadata); @@ -48,7 +48,10 @@ test("it should modify collection", async () => { test("it should store metadata", async () => { await chroma.reset(); - const collection = await chroma.createCollection({ name: "test", metadata: { test: "test" } }); + const collection = await chroma.createCollection({ + name: "test", + metadata: { test: "test" }, + }); expect(collection.metadata).toEqual({ test: "test" }); // get the collection diff --git a/clients/js/test/delete.collection.test.ts b/clients/js/test/delete.collection.test.ts index a192972b599..c4a3f8310ee 100644 --- a/clients/js/test/delete.collection.test.ts +++ b/clients/js/test/delete.collection.test.ts @@ -5,7 +5,11 @@ import { EMBEDDINGS, IDS, METADATAS } from "./data"; test("it should delete a collection", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + }); let count = await collection.count(); expect(count).toBe(3); var resp = await collection.delete({ where: { test: "test1" } }); @@ -14,6 +18,6 @@ test("it should delete a collection", async () => { var remainingEmbeddings = await collection.get(); expect(["test2", "test3"]).toEqual( - expect.arrayContaining(remainingEmbeddings.ids) + expect.arrayContaining(remainingEmbeddings.ids), ); }); diff --git a/clients/js/test/get.collection.test.ts b/clients/js/test/get.collection.test.ts index 4ff88d208ff..93323654e4c 100644 --- a/clients/js/test/get.collection.test.ts +++ b/clients/js/test/get.collection.test.ts @@ -5,7 +5,11 @@ import { DOCUMENTS, EMBEDDINGS, IDS, METADATAS } from "./data"; test("it should get a collection", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + }); const results = await collection.get({ ids: ["test1"] }); expect(results).toBeDefined(); expect(results).toBeInstanceOf(Object); @@ -23,12 +27,16 @@ test("it should get a collection", async () => { test("wrong code returns an error", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + }); const results = await collection.get({ where: { //@ts-ignore supposed to fail test: { $contains: "hello" }, - } + }, }); expect(results.error).toBeDefined(); expect(results.error).toContain("ValueError"); @@ -37,28 +45,56 @@ test("wrong code returns an error", async () => { test("it should get embedding with matching documents", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); - const results2 = await collection.get({ whereDocument: { $contains: "This is a test" } }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); + const results2 = await collection.get({ + whereDocument: { $contains: "This is a test" }, + }); expect(results2).toBeDefined(); expect(results2).toBeInstanceOf(Object); expect(results2.ids.length).toBe(1); expect(["test1"]).toEqual(expect.arrayContaining(results2.ids)); }); +test("it should get records not matching", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); + const results2 = await collection.get({ + whereDocument: { $not_contains: "This is another" }, + }); + expect(results2).toBeDefined(); + expect(results2).toBeInstanceOf(Object); + expect(results2.ids.length).toBe(2); + expect(["test1", "test3"]).toEqual(expect.arrayContaining(results2.ids)); +}); + test("test gt, lt, in a simple small way", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + }); const items = await collection.get({ where: { float_value: { $gt: -1.4 } } }); expect(items.ids.length).toBe(2); expect(["test2", "test3"]).toEqual(expect.arrayContaining(items.ids)); }); - test("it should throw an error if the collection does not exist", async () => { await chroma.reset(); await expect( - async () => await chroma.getCollection({ name: "test" }) + async () => await chroma.getCollection({ name: "test" }), ).rejects.toThrow(Error); }); diff --git a/clients/js/test/initClientWithAuth.ts b/clients/js/test/initClientWithAuth.ts index 0fd55c4e7d2..e92ad4df508 100644 --- a/clients/js/test/initClientWithAuth.ts +++ b/clients/js/test/initClientWithAuth.ts @@ -1,16 +1,34 @@ -import {ChromaClient} from "../src/ChromaClient"; +import { ChromaClient } from "../src/ChromaClient"; import { CloudClient } from "../src/CloudClient"; const PORT = process.env.PORT || "8000"; const URL = "http://localhost:" + PORT; -export const chromaBasic = new ChromaClient({path: URL, auth: {provider: "basic", credentials: "admin:admin"}}); -export const chromaTokenDefault = new ChromaClient({path: URL, auth: {provider: "token", credentials: "test-token"}}); +export const chromaBasic = new ChromaClient({ + path: URL, + auth: { provider: "basic", credentials: "admin:admin" }, +}); +export const chromaTokenDefault = new ChromaClient({ + path: URL, + auth: { provider: "token", credentials: "test-token" }, +}); export const chromaTokenBearer = new ChromaClient({ - path: URL, - auth: {provider: "token", credentials: "test-token", providerOptions: {headerType: "AUTHORIZATION"}} + path: URL, + auth: { + provider: "token", + credentials: "test-token", + providerOptions: { headerType: "AUTHORIZATION" }, + }, }); export const chromaTokenXToken = new ChromaClient({ - path: URL, - auth: {provider: "token", credentials: "test-token", providerOptions: {headerType: "X_CHROMA_TOKEN"}} + path: URL, + auth: { + provider: "token", + credentials: "test-token", + providerOptions: { headerType: "X_CHROMA_TOKEN" }, + }, +}); +export const cloudClient = new CloudClient({ + apiKey: "test-token", + cloudPort: PORT, + cloudHost: "http://localhost", }); -export const cloudClient = new CloudClient({apiKey: "test-token", cloudPort: PORT, cloudHost: "http://localhost"}) diff --git a/clients/js/test/query.collection.test.ts b/clients/js/test/query.collection.test.ts index 878ed0a71df..2809716535c 100644 --- a/clients/js/test/query.collection.test.ts +++ b/clients/js/test/query.collection.test.ts @@ -6,8 +6,7 @@ import { EMBEDDINGS, IDS, METADATAS, DOCUMENTS } from "./data"; import { IEmbeddingFunction } from "../src/embeddings/IEmbeddingFunction"; export class TestEmbeddingFunction implements IEmbeddingFunction { - - constructor() { } + constructor() {} public async generate(texts: string[]): Promise { let embeddings: number[][] = []; @@ -22,7 +21,10 @@ test("it should query a collection", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); await collection.add({ ids: IDS, embeddings: EMBEDDINGS }); - const results = await collection.query({ queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], nResults: 2 }); + const results = await collection.query({ + queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + nResults: 2, + }); expect(results).toBeDefined(); expect(results).toBeInstanceOf(Object); expect(["test1", "test2"]).toEqual(expect.arrayContaining(results.ids[0])); @@ -33,12 +35,17 @@ test("it should query a collection", async () => { test("it should get embedding with matching documents", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); const results = await collection.query({ queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], nResults: 3, - whereDocument: { $contains: "This is a test" } + whereDocument: { $contains: "This is a test" }, }); // it should only return doc1 @@ -48,14 +55,14 @@ test("it should get embedding with matching documents", async () => { expect(["test1"]).toEqual(expect.arrayContaining(results.ids[0])); expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0])); expect(["This is a test"]).toEqual( - expect.arrayContaining(results.documents[0]) + expect.arrayContaining(results.documents[0]), ); const results2 = await collection.query({ queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], nResults: 3, whereDocument: { $contains: "This is a test" }, - include: [IncludeEnum.Embeddings] + include: [IncludeEnum.Embeddings], }); // expect(results2.embeddings[0][0]).toBeInstanceOf(Array); @@ -63,18 +70,48 @@ test("it should get embedding with matching documents", async () => { expect(results2.embeddings![0][0]).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); }); +test("it should exclude documents matching - not_contains", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); + + const results = await collection.query({ + queryEmbeddings: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + nResults: 3, + whereDocument: { $not_contains: "This is a test" }, + }); + + // it should only return doc1 + expect(results).toBeDefined(); + expect(results).toBeInstanceOf(Object); + expect(results.ids.length).toBe(1); + expect(["test2", "test3"]).toEqual(expect.arrayContaining(results.ids[0])); +}); // test queryTexts test("it should query a collection with text", async () => { await chroma.reset(); let embeddingFunction = new TestEmbeddingFunction(); - const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embeddingFunction, + }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); const results = await collection.query({ queryTexts: ["test"], nResults: 3, - whereDocument: { $contains: "This is a test" } + whereDocument: { $contains: "This is a test" }, }); expect(results).toBeDefined(); @@ -83,21 +120,28 @@ test("it should query a collection with text", async () => { expect(["test1"]).toEqual(expect.arrayContaining(results.ids[0])); expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0])); expect(["This is a test"]).toEqual( - expect.arrayContaining(results.documents[0]) + expect.arrayContaining(results.documents[0]), ); -}) - +}); test("it should query a collection with text and where", async () => { await chroma.reset(); let embeddingFunction = new TestEmbeddingFunction(); - const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embeddingFunction, + }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); const results = await collection.query({ queryTexts: ["test"], nResults: 3, - where: { "float_value" : 2 } + where: { float_value: 2 }, }); expect(results).toBeDefined(); @@ -106,21 +150,28 @@ test("it should query a collection with text and where", async () => { expect(["test3"]).toEqual(expect.arrayContaining(results.ids[0])); expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0])); expect(["This is a third test"]).toEqual( - expect.arrayContaining(results.documents[0]) + expect.arrayContaining(results.documents[0]), ); -}) - +}); test("it should query a collection with text and where in", async () => { await chroma.reset(); let embeddingFunction = new TestEmbeddingFunction(); - const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embeddingFunction, + }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); const results = await collection.query({ queryTexts: ["test"], nResults: 3, - where: { "float_value" : { '$in': [2,5,10] }} + where: { float_value: { $in: [2, 5, 10] } }, }); expect(results).toBeDefined(); @@ -129,20 +180,28 @@ test("it should query a collection with text and where in", async () => { expect(["test3"]).toEqual(expect.arrayContaining(results.ids[0])); expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0])); expect(["This is a third test"]).toEqual( - expect.arrayContaining(results.documents[0]) + expect.arrayContaining(results.documents[0]), ); -}) +}); test("it should query a collection with text and where nin", async () => { await chroma.reset(); let embeddingFunction = new TestEmbeddingFunction(); - const collection = await chroma.createCollection({ name: "test", embeddingFunction: embeddingFunction }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); + const collection = await chroma.createCollection({ + name: "test", + embeddingFunction: embeddingFunction, + }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); const results = await collection.query({ queryTexts: ["test"], nResults: 3, - where: { "float_value" : { '$nin': [-2,0] }} + where: { float_value: { $nin: [-2, 0] } }, }); expect(results).toBeDefined(); @@ -151,6 +210,6 @@ test("it should query a collection with text and where nin", async () => { expect(["test3"]).toEqual(expect.arrayContaining(results.ids[0])); expect(["test2"]).not.toEqual(expect.arrayContaining(results.ids[0])); expect(["This is a third test"]).toEqual( - expect.arrayContaining(results.documents[0]) + expect.arrayContaining(results.documents[0]), ); -}) +}); diff --git a/clients/js/test/update.collection.test.ts b/clients/js/test/update.collection.test.ts index e96f21d5b06..77537ac6bfb 100644 --- a/clients/js/test/update.collection.test.ts +++ b/clients/js/test/update.collection.test.ts @@ -6,7 +6,12 @@ import { IDS, DOCUMENTS, EMBEDDINGS, METADATAS } from "./data"; test("it should get embedding with matching documents", async () => { await chroma.reset(); const collection = await chroma.createCollection({ name: "test" }); - await collection.add({ ids: IDS, embeddings: EMBEDDINGS, metadatas: METADATAS, documents: DOCUMENTS }); + await collection.add({ + ids: IDS, + embeddings: EMBEDDINGS, + metadatas: METADATAS, + documents: DOCUMENTS, + }); const results = await collection.get({ ids: ["test1"], @@ -14,7 +19,7 @@ test("it should get embedding with matching documents", async () => { IncludeEnum.Embeddings, IncludeEnum.Metadatas, IncludeEnum.Documents, - ] + ], }); expect(results).toBeDefined(); expect(results).toBeInstanceOf(Object); @@ -24,7 +29,7 @@ test("it should get embedding with matching documents", async () => { ids: ["test1"], embeddings: [[1, 2, 3, 4, 5, 6, 7, 8, 9, 11]], metadatas: [{ test: "test1new" }], - documents: ["doc1new"] + documents: ["doc1new"], }); const results2 = await collection.get({ @@ -33,7 +38,7 @@ test("it should get embedding with matching documents", async () => { IncludeEnum.Embeddings, IncludeEnum.Metadatas, IncludeEnum.Documents, - ] + ], }); expect(results2).toBeDefined(); expect(results2).toBeInstanceOf(Object); diff --git a/clients/js/test/upsert.collections.test.ts b/clients/js/test/upsert.collections.test.ts index 9ce00820e2d..0fc1cdf1ef4 100644 --- a/clients/js/test/upsert.collections.test.ts +++ b/clients/js/test/upsert.collections.test.ts @@ -1,27 +1,26 @@ -import { expect, test } from '@jest/globals'; -import chroma from './initClient' +import { expect, test } from "@jest/globals"; +import chroma from "./initClient"; +test("it should upsert embeddings to a collection", async () => { + await chroma.reset(); + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["test1", "test2"]; + const embeddings = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + ]; + await collection.add({ ids, embeddings }); + const count = await collection.count(); + expect(count).toBe(2); -test('it should upsert embeddings to a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection({ name: "test" }); - const ids = ['test1', 'test2'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add({ ids, embeddings }) - const count = await collection.count() - expect(count).toBe(2) + const ids2 = ["test2", "test3"]; + const embeddings2 = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 15], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + ]; - const ids2 = ["test2", "test3"] - const embeddings2 = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 15], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - ] + await collection.upsert({ ids: ids2, embeddings: embeddings2 }); - await collection.upsert({ ids: ids2, embeddings: embeddings2 }) - - const count2 = await collection.count() - expect(count2).toBe(3) -}) + const count2 = await collection.count(); + expect(count2).toBe(3); +}); diff --git a/clients/js/tsconfig.json b/clients/js/tsconfig.json index 632127ed709..c46b79fc3ba 100644 --- a/clients/js/tsconfig.json +++ b/clients/js/tsconfig.json @@ -1,7 +1,5 @@ { - "include": [ - "src" - ], + "include": ["src"], "compilerOptions": { "declaration": true, "module": "ESNext", diff --git a/clients/js/tsup.config.ts b/clients/js/tsup.config.ts index 4b0cd8e264e..aff2492dbea 100644 --- a/clients/js/tsup.config.ts +++ b/clients/js/tsup.config.ts @@ -1,32 +1,32 @@ -import { defineConfig, Options } from 'tsup' -import fs from 'fs' +import { defineConfig, Options } from "tsup"; +import fs from "fs"; export default defineConfig((options: Options) => { const commonOptions: Partial = { entry: { - chromadb: 'src/index.ts' + chromadb: "src/index.ts", }, sourcemap: true, dts: true, - ...options - } + ...options, + }; return [ { ...commonOptions, - format: ['esm'], - outExtension: () => ({ js: '.mjs' }), + format: ["esm"], + outExtension: () => ({ js: ".mjs" }), clean: true, async onSuccess() { // Support Webpack 4 by pointing `"module"` to a file with a `.js` extension - fs.copyFileSync('dist/chromadb.mjs', 'dist/chromadb.legacy-esm.js') - } + fs.copyFileSync("dist/chromadb.mjs", "dist/chromadb.legacy-esm.js"); + }, }, { ...commonOptions, - format: 'cjs', - outDir: './dist/cjs/', - outExtension: () => ({ js: '.cjs' }) - } - ] -}) + format: "cjs", + outDir: "./dist/cjs/", + outExtension: () => ({ js: ".cjs" }), + }, + ]; +}); diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index b62c002d095..edd0c00d7cf 100644 --- a/clients/python/pyproject.toml +++ b/clients/python/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ 'typing_extensions >= 4.5.0', 'tenacity>=8.2.3', 'PyYAML>=6.0.0', + 'orjson>=3.9.12', ] [tool.black] diff --git a/clients/python/requirements.txt b/clients/python/requirements.txt index 1242bf7d7e0..b977b03f064 100644 --- a/clients/python/requirements.txt +++ b/clients/python/requirements.txt @@ -9,3 +9,4 @@ PyYAML>=6.0.0 requests >= 2.28 tenacity>=8.2.3 typing_extensions >= 4.5.0 +orjson>=3.9.12 diff --git a/docker-compose.test-auth.yml b/docker-compose.test-auth.yml index 259d4c54e79..d3297b5a04f 100644 --- a/docker-compose.test-auth.yml +++ b/docker-compose.test-auth.yml @@ -11,7 +11,7 @@ services: dockerfile: Dockerfile volumes: - chroma-data:/chroma/chroma - command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30 + command: "--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30" environment: - ANONYMIZED_TELEMETRY=False - ALLOW_RESET=True diff --git a/docker-compose.test.yml b/docker-compose.test.yml index c8cae63b3eb..4384bad1982 100644 --- a/docker-compose.test.yml +++ b/docker-compose.test.yml @@ -11,7 +11,7 @@ services: dockerfile: Dockerfile volumes: - chroma-data:/chroma/chroma - command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30 + command: "--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30" environment: - ANONYMIZED_TELEMETRY=False - ALLOW_RESET=True diff --git a/docker-compose.yml b/docker-compose.yml index 3e023109458..2d34395ce59 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,21 +15,29 @@ services: # Default configuration for persist_directory in chromadb/config.py # Read more about deployments: https://docs.trychroma.com/deployment - chroma-data:/chroma/chroma - command: uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30 + command: "--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30" environment: - IS_PERSISTENT=TRUE - CHROMA_SERVER_AUTH_PROVIDER=${CHROMA_SERVER_AUTH_PROVIDER} - CHROMA_SERVER_AUTH_CREDENTIALS_FILE=${CHROMA_SERVER_AUTH_CREDENTIALS_FILE} - CHROMA_SERVER_AUTH_CREDENTIALS=${CHROMA_SERVER_AUTH_CREDENTIALS} - CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=${CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER} + - CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER=${CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER} - PERSIST_DIRECTORY=${PERSIST_DIRECTORY:-/chroma/chroma} - CHROMA_OTEL_EXPORTER_ENDPOINT=${CHROMA_OTEL_EXPORTER_ENDPOINT} - CHROMA_OTEL_EXPORTER_HEADERS=${CHROMA_OTEL_EXPORTER_HEADERS} - CHROMA_OTEL_SERVICE_NAME=${CHROMA_OTEL_SERVICE_NAME} - CHROMA_OTEL_GRANULARITY=${CHROMA_OTEL_GRANULARITY} - CHROMA_SERVER_NOFILE=${CHROMA_SERVER_NOFILE} + restart: unless-stopped # possible values are: "no", always", "on-failure", "unless-stopped" ports: - - 8000:8000 + - "8000:8000" + healthcheck: + # Adjust below to match your container port + test: [ "CMD", "curl", "-f", "http://localhost:8000/api/v1/heartbeat" ] + interval: 30s + timeout: 10s + retries: 3 networks: - net diff --git a/docs/cip/CIP-01022024_SSL_Verify_Client_Config.md b/docs/cip/CIP-01022024_SSL_Verify_Client_Config.md new file mode 100644 index 00000000000..2448af11c88 --- /dev/null +++ b/docs/cip/CIP-01022024_SSL_Verify_Client_Config.md @@ -0,0 +1,68 @@ +# CIP-01022024 SSL Verify Client Config + +## Status + +Current Status: `Under Discussion` + +## Motivation + +The motivation for this change is to enhance security and flexibility in Chroma's client API. Users need the ability to +configure SSL contexts to trust custom CA certificates or self-signed certificates, which is not straightforward with +the current setup. This capability is crucial for organizations that operate their own CA or for developers who need to +test their applications in environments where certificates from a recognized CA are not available or practical. + +The suggested change entails a server-side certificate be available, but this CIP does not prescribe how such +certificate should be configured or obtained. In our testing, we used a self-signed certificate generated with +`openssl` and configured the client to trust the certificate. We also experiment with a SSL-terminated proxy server. +Both of approaches yielded the same results. + +> **IMPORTANT:** It should be noted that we do not recommend or encourage the use of self-signed certificates in +> production environments. + +We also provide a sample notebook that to help the reader run a local Chroma server with a self-signed certificate and +configure the client to trust the certificate. The notebook can be found +in [assets/CIP-01022024-test_self_signed.ipynb](./assets/CIP-01022024-test_self_signed.ipynb). + +## Public Interfaces + +> **Note:** The following changes are only applicable to Chroma HttpClient. + +New settings variable `chroma_server_ssl_verify` accepting either a boolean or a path to a certificate file. If the +value is a path to a certificate file, the file will be used to verify the server's certificate. If the value is a +boolean, the SSL certificate verification can be bypassed (`false`) or enforced (`true`). + +The value is passed as `verify` parameter to `requests.Session` of the `FastAPI` client. See +requests [documentation](https://requests.readthedocs.io/en/latest/user/advanced/#ssl-cert-verification) for +more details. + +Example Usage: + +```python +import chromadb +from chromadb import Settings +client = chromadb.HttpClient(host="localhost",port="8443",ssl=True, settings=Settings(chroma_server_ssl_verify='./servercert.pem')) +# or with boolean +client = chromadb.HttpClient(host="localhost",port="8443",ssl=True, settings=Settings(chroma_server_ssl_verify=False)) +``` + +### Resources + +- https://requests.readthedocs.io/en/latest/api/#requests.request +- https://www.geeksforgeeks.org/ssl-certificate-verification-python-requests/ + +## Proposed Changes + +The proposed changes are mentioned in the public interfaces. + +## Compatibility, Deprecation, and Migration Plan + +The change is not backward compatible from client's perspective as the lack of the feature in prior clients will cause +an error when passing the new settings parameter. Server-side is not affected by this change. + +## Test Plan + +API tests with SSL verification enabled and a self-signed certificate. + +## Rejected Alternatives + +N/A diff --git a/docs/cip/assets/CIP-01022024-test_self_signed.ipynb b/docs/cip/assets/CIP-01022024-test_self_signed.ipynb new file mode 100644 index 00000000000..d607b51824b --- /dev/null +++ b/docs/cip/assets/CIP-01022024-test_self_signed.ipynb @@ -0,0 +1,119 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "source": [ + "# Generate a Certificate\n", + "\n", + "```bash\n", + "openssl req -new -newkey rsa:2048 -sha256 -days 365 -nodes -x509 \\\n", + " -keyout ./serverkey.pem \\\n", + " -out ./servercert.pem \\\n", + " -subj \"/O=Chroma/C=US\" \\\n", + " -config chromadb/test/openssl.cnf\n", + "```\n", + "\n", + "> Note: The above command should be executed at the root of the repo (openssl.cnf uses relative path)\n" + ], + "metadata": { + "collapsed": false + }, + "id": "faa8cefb6825fe83" + }, + { + "cell_type": "markdown", + "source": [ + "# Start the server\n", + "\n", + "```bash\n", + "uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8443 \\\n", + " --proxy-headers --log-config chromadb/log_config.yml --ssl-keyfile ./serverkey.pem --ssl-certfile ./servercert.pem\n", + "```" + ], + "metadata": { + "collapsed": false + }, + "id": "e084285e11c3747d" + }, + { + "cell_type": "markdown", + "source": [ + "# Test with cert as SSL verify string" + ], + "metadata": { + "collapsed": false + }, + "id": "130df9c0a6d67b52" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "initial_id", + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "from chromadb import Settings\n", + "import chromadb\n", + "client = chromadb.HttpClient(host=\"localhost\",port=\"8443\",ssl=True, settings=Settings(chroma_server_ssl_verify='./servercert.pem'))\n", + "print(client.heartbeat())" + ] + }, + { + "cell_type": "markdown", + "source": [ + "# Test with cert as SSL verify boolean" + ], + "metadata": { + "collapsed": false + }, + "id": "8223d0100df06ec4" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from chromadb import Settings\n", + "import chromadb\n", + "client = chromadb.HttpClient(host=\"localhost\",port=\"8443\",ssl=True, settings=Settings(chroma_server_ssl_verify=False))\n", + "print(client.heartbeat())" + ], + "metadata": { + "collapsed": false + }, + "id": "f7cf299721741c1", + "execution_count": null + }, + { + "cell_type": "code", + "outputs": [], + "source": [], + "metadata": { + "collapsed": false + }, + "id": "6231ac2ac38383c2" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/use_with/ollama.md b/examples/use_with/ollama.md new file mode 100644 index 00000000000..7b6977fec7d --- /dev/null +++ b/examples/use_with/ollama.md @@ -0,0 +1,40 @@ +# Ollama + +First let's run a local docker container with Ollama. We'll pull `nomic-embed-text` model: + +```bash +docker run -d -v ./ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama +docker exec -it ollama ollama run nomic-embed-text # press Ctrl+D to exit after model downloads successfully +# test it +curl http://localhost:11434/api/embeddings -d '{"model": "nomic-embed-text","prompt": "Here is an article about llamas..."}' +``` + +Now let's configure our OllamaEmbeddingFunction Embedding (python) function with the default Ollama endpoint: + +```python +import chromadb +from chromadb.utils.embedding_functions import OllamaEmbeddingFunction + +client = chromadb.PersistentClient(path="ollama") + +# create EF with custom endpoint +ef = OllamaEmbeddingFunction( + model_name="nomic-embed-text", + url="http://127.0.0.1:11434/api/embeddings", +) + +print(ef(["Here is an article about llamas..."])) +``` + +For JS users, you can use the `OllamaEmbeddingFunction` class to create embeddings: + +```javascript +const {OllamaEmbeddingFunction} = require('chromadb'); +const embedder = new OllamaEmbeddingFunction({ + url: "http://127.0.0.1:11434/api/embeddings", + model: "llama2" +}) + +// use directly +const embeddings = embedder.generate(["Here is an article about llamas..."]) +``` diff --git a/examples/use_with/roboflow/embeddings.ipynb b/examples/use_with/roboflow/embeddings.ipynb new file mode 100644 index 00000000000..b5d03d445a9 --- /dev/null +++ b/examples/use_with/roboflow/embeddings.ipynb @@ -0,0 +1,258 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "wFnbuH7qqotV" + }, + "source": [ + "# Use Roboflow with Chroma\n", + "\n", + "With [Roboflow Inference](https://inference.roboflow.com), you can calculate image embeddings using CLIP, a popular multimodal embedding model. You can then store these embeddings in Chroma for use in your application.\n", + "\n", + "In this guide, we are going to discuss how to load image embeddings into Chroma. We will discuss:\n", + "\n", + "1. How to set up Roboflow Inference\n", + "2. How to create a Chroma vector database\n", + "3. How to calculate CLIP embeddings with Inference\n", + "4. How to run a search query with Chroma\n", + "\n", + "## What is Roboflow Inference?\n", + "\n", + "[Roboflow Inference](https://inference.roboflow.com) is a scalable server through which you can run fine-tuned object detection, segmentation, and classification models, as well as popular foundation models such as CLIP.\n", + "\n", + "Inference handles all of the complexity associated with running vision models, from managing dependencies to maintaining your environment.\n", + "\n", + "Inference is trusted by enterprises around the world to manage vision models, with the hosted version powering millions of API calls each month.\n", + "\n", + "Inference runs in Docker and provides a HTTP interface through which to retrieve predictions.\n", + "\n", + "We will use Inference to calculate CLIP embeddings for our application.\n", + "\n", + "There are two ways to use Inference:\n", + "\n", + "1. On your device\n", + "2. Through the Inference API hosted by Roboflow\n", + "\n", + "In this guide, we will use the hosted Inference API." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "coXj8QiRrXfw" + }, + "source": [ + "### Step #1: Create a Chroma Vector Database\n", + "\n", + "To load and save image embeddings into Chroma, we first need images to embed. In this guide, we are going to use the COCO 128 dataset, a collection of 128 images from the Microsoft COCO dataset. This dataset is available on Roboflow Universe, a community that has shared more than 250,000 public computer vision datasets.\n", + "\n", + "To download the dataset, visit the COCO 128 web page, click “Download Dataset” and click \"show download code\" to get a download code:\n", + "\n", + "![COCO 128 dataset](https://media.roboflow.com/coco128.png)\n", + "\n", + "Here is the download code for the COCO 128 dataset:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4MboNZCZsTfK" + }, + "outputs": [], + "source": [ + "!pip install roboflow -q\n", + "\n", + "API_KEY = \"\"\n", + "\n", + "from roboflow import Roboflow\n", + "\n", + "rf = Roboflow(api_key=API_KEY)\n", + "project = rf.workspace(\"team-roboflow\").project(\"coco-128\")\n", + "dataset = project.version(2).download(\"yolov8\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "pf3aKIsGsTKD" + }, + "source": [ + "\n", + "Above, replace the value associated with the `API_KEY` variable with your Roboflow API key. [Learn how to retrieve your Robflow API key](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key).\n", + "\n", + "Now that we have a dataset ready, we can create a vector database and start loading embeddings.\n", + "\n", + "Install the Chroma Python client and supervision, which we will use to open images in this notebook, with the following command:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "smDcsb16rZdP" + }, + "outputs": [], + "source": [ + "!pip install chromadb supervision -q" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-IvKMl9IrcOJ" + }, + "source": [ + "Then, run the code below to calculate CLIP vectors for images in your dataset:" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "id": "BPN4-uvLrbhQ" + }, + "outputs": [], + "source": [ + "import chromadb\n", + "import os\n", + "from chromadb.utils.data_loaders import ImageLoader\n", + "from chromadb.utils.embedding_functions import RoboflowEmbeddingFunction\n", + "import uuid\n", + "import cv2\n", + "import supervision as sv\n", + "\n", + "SERVER_URL = \"https://infer.roboflow.com\"\n", + "\n", + "ef = RoboflowEmbeddingFunction(API_KEY, api_url = SERVER_URL)\n", + "\n", + "client = chromadb.PersistentClient(path=\"database\")\n", + "\n", + "data_loader = ImageLoader()\n", + "\n", + "collection = client.create_collection(name=\"images_db2\", embedding_function=ef, data_loader=data_loader, metadata={\"hnsw:space\": \"cosine\"})" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "gAFEc5FJu7oj", + "outputId": "4bf5d5b8-0c88-4ff4-bb83-dcf47178c770" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "['/content/COCO-128-2/train/images/000000000643_jpg.rf.fd058cb12cbda17a08a6254751aad243.jpg', '/content/COCO-128-2/train/images/000000000446_jpg.rf.4d5fe626e32b8b40408f9b711a10f04a.jpg', '/content/COCO-128-2/train/images/000000000321_jpg.rf.012f28b6a17e876bf2da17cab227c4cc.jpg', '/content/COCO-128-2/train/images/000000000315_jpg.rf.3613aa9bc01121a7949ecde10452ceef.jpg', '/content/COCO-128-2/train/images/000000000472_jpg.rf.68cb82fda7d6d5abec9e462eb191271a.jpg', '/content/COCO-128-2/train/images/000000000389_jpg.rf.1dd437dbb518e480b21652cf552e5b2d.jpg', '/content/COCO-128-2/train/images/000000000208_jpg.rf.5247f6d89d634629bfa005d8792613b4.jpg', '/content/COCO-128-2/train/images/000000000597_jpg.rf.2c8f04559f193762dc844986f6d60cad.jpg', '/content/COCO-128-2/train/images/000000000508_jpg.rf.4d5b0d34a3ddacbbae66a6d9a0c4daf5.jpg', '/content/COCO-128-2/train/images/000000000431_jpg.rf.56ab0d462037c5b588d11c7eb5b34278.jpg', '/content/COCO-128-2/train/images/000000000357_jpg.rf.fd60a5947f0f1b2ad6273d8ff87b6282.jpg', '/content/COCO-128-2/train/images/000000000138_jpg.rf.af439ef1c55dd8a4e4b142d186b9c957.jpg', '/content/COCO-128-2/train/images/000000000438_jpg.rf.c410916a61676ab359be5aebf293c85f.jpg', '/content/COCO-128-2/train/images/000000000419_jpg.rf.b58c3291ae18177c82e19195fd533614.jpg', '/content/COCO-128-2/train/images/000000000510_jpg.rf.a6383d3c4bffc15986bababf90bb2076.jpg', '/content/COCO-128-2/train/images/000000000143_jpg.rf.99b5580faf8d1ff7b0ac0b4345a4cf1a.jpg', '/content/COCO-128-2/train/images/000000000260_jpg.rf.be309dc699a9462430efc6dc624a0681.jpg', '/content/COCO-128-2/train/images/000000000110_jpg.rf.150f5495ca5c7842cfcd42a5aeae841f.jpg', '/content/COCO-128-2/train/images/000000000562_jpg.rf.9bb940c3cf544f101e88cd9052ddc3e0.jpg', '/content/COCO-128-2/train/images/000000000308_jpg.rf.220721ec7c9e25fe8596c5c64cc84a2c.jpg', '/content/COCO-128-2/train/images/000000000165_jpg.rf.eae14d5509bf0c9ceccddbb53a5f0c66.jpg', '/content/COCO-128-2/train/images/000000000415_jpg.rf.670afeffb0d21fd977df575a7a826b39.jpg', '/content/COCO-128-2/train/images/000000000436_jpg.rf.f75c349fa1c6f78054991af5238d8e0a.jpg', '/content/COCO-128-2/train/images/000000000192_jpg.rf.ad225cb1bc09bfe56f6282c3f7ce56af.jpg', '/content/COCO-128-2/train/images/000000000042_jpg.rf.a1f5c146b4b81881b19f342a148c0f51.jpg', '/content/COCO-128-2/train/images/000000000650_jpg.rf.1b74ba165c5a3513a3211d4a80b69e1c.jpg', '/content/COCO-128-2/train/images/000000000326_jpg.rf.02c19837e58093adec35e737066bb9ae.jpg', '/content/COCO-128-2/train/images/000000000149_jpg.rf.0a861d05b36be7927ab205acb325f4ce.jpg', '/content/COCO-128-2/train/images/000000000488_jpg.rf.c01187aceb8a49b901abe79739d6acdf.jpg', '/content/COCO-128-2/train/images/000000000241_jpg.rf.883e3e7bef174603fd5d2bbdbaf3a4ba.jpg', '/content/COCO-128-2/train/images/000000000370_jpg.rf.b13777ecf61d3edb14a6724a6331d4b7.jpg', '/content/COCO-128-2/train/images/000000000136_jpg.rf.71b51a4103c3e797f62a52a9d20fddfe.jpg', '/content/COCO-128-2/train/images/000000000625_jpg.rf.ce871c39393fefd9fd8671806761a1c8.jpg', '/content/COCO-128-2/train/images/000000000400_jpg.rf.fef965b3dde5237dcf1396f68c3b2a52.jpg', '/content/COCO-128-2/train/images/000000000073_jpg.rf.eb8e88e2239ba953b553f4e60b5d567c.jpg', '/content/COCO-128-2/train/images/000000000531_jpg.rf.5a9928283716bf2aac47963ca1a19afd.jpg', '/content/COCO-128-2/train/images/000000000612_jpg.rf.656879428df938a1a000bc255a193ccd.jpg', '/content/COCO-128-2/train/images/000000000142_jpg.rf.5d34f341f09bef6870b506337bb426ad.jpg', '/content/COCO-128-2/train/images/000000000036_jpg.rf.af0418c203165d3ee53b7dee5fe8b301.jpg', '/content/COCO-128-2/train/images/000000000089_jpg.rf.b82e9aa4a74633f95da698d712065b7a.jpg', '/content/COCO-128-2/train/images/000000000474_jpg.rf.bbd3bc1d1951dbca328a9a84c330b337.jpg', '/content/COCO-128-2/train/images/000000000404_jpg.rf.ca3ae1b40e22e5f6044ad9606b069ed7.jpg', '/content/COCO-128-2/train/images/000000000154_jpg.rf.300698916140dd41f6fda1c194d7b00d.jpg', '/content/COCO-128-2/train/images/000000000133_jpg.rf.8a7d1da21b04545e5543d28373d75cf1.jpg', '/content/COCO-128-2/train/images/000000000071_jpg.rf.1e9a11ad22ba304c5519e32b3ab47a48.jpg', '/content/COCO-128-2/train/images/000000000196_jpg.rf.8e48b2a4a9bd63fcd02ad61708d18ef2.jpg', '/content/COCO-128-2/train/images/000000000382_jpg.rf.c172a50ccf4da06a423497fac9d12579.jpg', '/content/COCO-128-2/train/images/000000000086_jpg.rf.556402c74eccf93483195dc2000ceeb2.jpg', '/content/COCO-128-2/train/images/000000000636_jpg.rf.e4dbb1e2ce37580cbc890709cad177a7.jpg', '/content/COCO-128-2/train/images/000000000328_jpg.rf.13cb1874b4f817acce9a8ad106b4225b.jpg', '/content/COCO-128-2/train/images/000000000113_jpg.rf.0b3da2103c183dbd0d3ae9a364b48458.jpg', '/content/COCO-128-2/train/images/000000000127_jpg.rf.b1a2420c07d63c4415880a90e639579f.jpg', '/content/COCO-128-2/train/images/000000000623_jpg.rf.85a92f3af326182db18c9f891a2a2d88.jpg', '/content/COCO-128-2/train/images/000000000532_jpg.rf.7f786bd3fd5dce29c2a1c58cd32b68d1.jpg', '/content/COCO-128-2/train/images/000000000338_jpg.rf.db620506b174e29042f3758665d1b384.jpg', '/content/COCO-128-2/train/images/000000000397_jpg.rf.9a764f8819f8ef9c3f93fdf33e9f4929.jpg', '/content/COCO-128-2/train/images/000000000294_jpg.rf.f4aaaee71c6e53cb4af0d42b7bc64c6b.jpg', '/content/COCO-128-2/train/images/000000000629_jpg.rf.d678871351924516e2310e8cbc18feeb.jpg', '/content/COCO-128-2/train/images/000000000250_jpg.rf.fe4daef71d5cf0dedf2d70ec1a053e05.jpg', '/content/COCO-128-2/train/images/000000000025_jpg.rf.782fe78a513b7eeded6172306f4f502c.jpg', '/content/COCO-128-2/train/images/000000000395_jpg.rf.7b3a9f039340b28eeea8eca51e875130.jpg', '/content/COCO-128-2/train/images/000000000540_jpg.rf.d42cc5cec9a137294c1d0dd81cacceaf.jpg', '/content/COCO-128-2/train/images/000000000061_jpg.rf.6945c178a14b90b01cd7d8b1e60e6a22.jpg', '/content/COCO-128-2/train/images/000000000387_jpg.rf.f6d68f4e3c90d097157a7d8fe1839b34.jpg', '/content/COCO-128-2/train/images/000000000529_jpg.rf.16dbfc2ac51287dba01ca3ae5028ece0.jpg', '/content/COCO-128-2/train/images/000000000109_jpg.rf.a74c9fd0bdb672415629919c2c802f2d.jpg', '/content/COCO-128-2/train/images/000000000368_jpg.rf.54fa547d31097e00d82aebd06a36df59.jpg', '/content/COCO-128-2/train/images/000000000620_jpg.rf.1e47693b5bf354d0ea9048db0be8227e.jpg', '/content/COCO-128-2/train/images/000000000283_jpg.rf.5b072b49882659b6ec9c9454d9623390.jpg', '/content/COCO-128-2/train/images/000000000572_jpg.rf.42b48081afe86d9e293734b0da064a71.jpg', '/content/COCO-128-2/train/images/000000000349_jpg.rf.9ed53bdff4c93ccc59fbe0b69de28cb5.jpg', '/content/COCO-128-2/train/images/000000000471_jpg.rf.faa1965b86263f4b92754c0495695c7e.jpg', '/content/COCO-128-2/train/images/000000000151_jpg.rf.f17fad0b1976f3f2dd638666617c159c.jpg', '/content/COCO-128-2/train/images/000000000312_jpg.rf.07ecf1a16342778b938ec2add28a7169.jpg', '/content/COCO-128-2/train/images/000000000443_jpg.rf.81f83991ba79c61e94912b2d34699024.jpg', '/content/COCO-128-2/train/images/000000000569_jpg.rf.af8d85de4a25be55832e3044098ea670.jpg', '/content/COCO-128-2/train/images/000000000634_jpg.rf.2feb5ee0e764217c6796acd17da1b7fa.jpg', '/content/COCO-128-2/train/images/000000000077_jpg.rf.520601990c512b7983888cbedfe14ca1.jpg', '/content/COCO-128-2/train/images/000000000450_jpg.rf.2cec6f80e01101afeb2060b0aac48e99.jpg', '/content/COCO-128-2/train/images/000000000194_jpg.rf.b6838f3b8a3e8626f1229fa694132bb9.jpg', '/content/COCO-128-2/train/images/000000000081_jpg.rf.5ac5126a29a5565691c27016453cb17b.jpg', '/content/COCO-128-2/train/images/000000000575_jpg.rf.6d3a2aa58cc3bd4db498d1a7b7a2c9d4.jpg', '/content/COCO-128-2/train/images/000000000428_jpg.rf.b5e82e0c67b750a2ac3eb0168d2b18d4.jpg', '/content/COCO-128-2/train/images/000000000009_jpg.rf.856f80d728927e943a5bccfdf49dd677.jpg', '/content/COCO-128-2/train/images/000000000641_jpg.rf.2ad873d1d358c17ad3149fac98f1e4bf.jpg', '/content/COCO-128-2/train/images/000000000257_jpg.rf.f6550733ae637214e209517d35c363b0.jpg', '/content/COCO-128-2/train/images/000000000064_jpg.rf.01bec3a770d2bc9cc8eb02d41a0c7f79.jpg', '/content/COCO-128-2/train/images/000000000360_jpg.rf.625cc871e851ded6de7bfe7687760f35.jpg', '/content/COCO-128-2/train/images/000000000094_jpg.rf.df1e8da2f564e0dc1f3e6401c05a1481.jpg', '/content/COCO-128-2/train/images/000000000309_jpg.rf.db7b22492fbb7f8d205fd1d00fe2280a.jpg', '/content/COCO-128-2/train/images/000000000030_jpg.rf.15401567b6ff1d3787995bde6eeee471.jpg', '/content/COCO-128-2/train/images/000000000584_jpg.rf.65a9bac7029d5afffc477baa9d3b43bc.jpg', '/content/COCO-128-2/train/images/000000000595_jpg.rf.8aa06812ef201b3ee0078533b0bbbdee.jpg', '/content/COCO-128-2/train/images/000000000564_jpg.rf.4ec8ed4abf0997c8cb5cea298fef3465.jpg', '/content/COCO-128-2/train/images/000000000307_jpg.rf.237662256b060aa457a2a09d95609b32.jpg', '/content/COCO-128-2/train/images/000000000560_jpg.rf.26e523d9666e9eadffd6df0f9c6754f0.jpg', '/content/COCO-128-2/train/images/000000000359_jpg.rf.01e0e6dc67e6fe503c7939d890bcd6d6.jpg', '/content/COCO-128-2/train/images/000000000263_jpg.rf.25c52657eff0a27882c0d1de8c72fe75.jpg', '/content/COCO-128-2/train/images/000000000605_jpg.rf.da1c1eeb4b1bca2c5862b31c0221a9b5.jpg', '/content/COCO-128-2/train/images/000000000514_jpg.rf.b65ef17c163528ff0f3c2131fd346a51.jpg', '/content/COCO-128-2/train/images/000000000486_jpg.rf.407a3acfe23dbf206d95b58a5d0aea37.jpg', '/content/COCO-128-2/train/images/000000000384_jpg.rf.616518e6584f3a3b798d270382dd2b16.jpg', '/content/COCO-128-2/train/images/000000000542_jpg.rf.f94e117d2fb63254d26cf370bcd0e86f.jpg', '/content/COCO-128-2/train/images/000000000201_jpg.rf.602495fdc61e3930f33684db99b1a869.jpg', '/content/COCO-128-2/train/images/000000000490_jpg.rf.b4d25b276034f3a0627399220c2cc0b8.jpg', '/content/COCO-128-2/train/images/000000000536_jpg.rf.04e1e3b47f96c565c19c6593aed44119.jpg', '/content/COCO-128-2/train/images/000000000332_jpg.rf.72bbbb24196a6387cbf5618e11f5f2ce.jpg', '/content/COCO-128-2/train/images/000000000074_jpg.rf.98fe5630d85530144d222acb65f55f15.jpg', '/content/COCO-128-2/train/images/000000000599_jpg.rf.06b89323b92ef324fdd5e7ecf7f20b9b.jpg', '/content/COCO-128-2/train/images/000000000581_jpg.rf.d56a78fa63c89e49f8b7092379a45c67.jpg', '/content/COCO-128-2/train/images/000000000049_jpg.rf.bfa6e0ac33a75f530011d8b3b50a3b5c.jpg', '/content/COCO-128-2/train/images/000000000459_jpg.rf.b6b7189271d5e7a826b5dad0923fbb87.jpg', '/content/COCO-128-2/train/images/000000000144_jpg.rf.09821163c359a738ad86e72865a310dc.jpg', '/content/COCO-128-2/train/images/000000000589_jpg.rf.a7b35721e40dfa21973cd1ab02a8fab4.jpg', '/content/COCO-128-2/train/images/000000000034_jpg.rf.a33e87d94b16c1112e8d9946fee784b9.jpg', '/content/COCO-128-2/train/images/000000000078_jpg.rf.afbc984af561c845df5edc5f96c4037f.jpg', '/content/COCO-128-2/train/images/000000000164_jpg.rf.276e7a583c850858df0541c60a19f28c.jpg', '/content/COCO-128-2/train/images/000000000092_jpg.rf.69d6172284e6afd6017db6ea911213ca.jpg', '/content/COCO-128-2/train/images/000000000322_jpg.rf.ca66cb6f129e0893f5f9139c3ec1617b.jpg', '/content/COCO-128-2/train/images/000000000544_jpg.rf.2762a5e736ff95a5e68bad19b86210e6.jpg', '/content/COCO-128-2/train/images/000000000072_jpg.rf.244aee5e871d00c83983a224a5eb8ed5.jpg', '/content/COCO-128-2/train/images/000000000590_jpg.rf.7d55377855a996a3c6ee00e8d98307b7.jpg', '/content/COCO-128-2/train/images/000000000502_jpg.rf.997919f6f6bc7d0929387ffcb6a49b24.jpg', '/content/COCO-128-2/train/images/000000000626_jpg.rf.b776e538820040987b54f49d4a2bb7fb.jpg', '/content/COCO-128-2/train/images/000000000247_jpg.rf.892bc95842cc54b7e6972677db7c5928.jpg', '/content/COCO-128-2/train/images/000000000491_jpg.rf.dd6cde4b463637c688bcae9dbb0488f0.jpg', '/content/COCO-128-2/train/images/000000000520_jpg.rf.30218d72a576772917cc478208e40924.jpg', '/content/COCO-128-2/train/images/000000000394_jpg.rf.757dedbde46dc32eed390faa679fc4f5.jpg']\n" + ] + } + ], + "source": [ + "IMAGE_DIR = dataset.location + \"/train/images\"\n", + "\n", + "documents = [os.path.join(IMAGE_DIR, img) for img in os.listdir(IMAGE_DIR)]\n", + "uris = [os.path.join(IMAGE_DIR, img) for img in os.listdir(IMAGE_DIR)]\n", + "ids = [str(uuid.uuid4()) for _ in range(len(documents))]\n", + "\n", + "collection.add(\n", + " uris=uris,\n", + " ids=ids,\n", + " metadatas=[{\"file\": file} for file in documents]\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "y-ai8v7BrozZ" + }, + "source": [ + "If you have downloaded custom images from a source other than the Roboflow snippet earlier in this notebook, replace `IMAGE_DIR` with the folder where your images are stored.\n", + "\n", + "In this code snippet, we create a new Chroma database called `images`. Our database will use cosine similarity for embedding comparisons.\n", + "\n", + "We calculate CLIP embeddings for all images in the `COCO128/train/images` folder using Inference. We save the embeddings in Chroma using the `collection.add()` method.\n", + "\n", + "We store the file names associated with each image in the `documents` variable, and embeddings in `embeddings`.\n", + "\n", + "If you want to use the hosted version of Roboflow Inference to calculate embeddings, replace the `SERVER_URL` value with `https://infer.roboflow.com`. We use the RoboflowEmbeddingFunction, built in to Chroma, to interact with Inference.\n", + "\n", + "Run the script above to calculate embeddings for a folder of images and save them in your database.\n", + "\n", + "We now have a vector database that contains some embeddings. Great! Let’s move on to the fun part: running a search query on our database." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "KCWbrXsbrpmI" + }, + "source": [ + "### Step #3: Run a Search Query\n", + "\n", + "To run a search query, we need a text embedding of a query. For example, if we want to find vegetables in our collection of 128 images from the COCO dataset, we need to have a text embedding for the search phrase “baseball”.\n", + "\n", + "To calculate a text embedding, we can use Inference through the embedding function we defined earlier:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4DeO6T7xrs2K" + }, + "outputs": [], + "source": [ + "query = \"baseball\"\n", + "\n", + "results = collection.query(\n", + " n_results=3,\n", + " query_texts=query\n", + ")\n", + "top_result = results[\"metadatas\"][0][0][\"file\"]\n", + "\n", + "sv.plot_image(cv2.imread(top_result))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Me2eRcYPrrXL" + }, + "source": [ + "Our code returns the name of the image with the most similar embedding to the embedding of our text query.\n", + "\n", + "The top result is an image of a child holding a baseball glove in a park. Chroma successfully returned an image that matched our prompt." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.4" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file diff --git a/go/Dockerfile b/go/Dockerfile new file mode 100644 index 00000000000..ebc9dadb335 --- /dev/null +++ b/go/Dockerfile @@ -0,0 +1,19 @@ +FROM golang:bookworm as builder +WORKDIR /build-dir +RUN apt-get update && apt-get install -y make git bash + +ADD ./go/go.mod ./go.mod +ADD ./go/go.sum ./go.sum +RUN go mod download + +ADD ./go/ ./ +ENV GOCACHE=/root/.cache/go-build +RUN --mount=type=cache,target="/root/.cache/go-build" make + +FROM debian:bookworm-slim + +COPY --from=builder /build-dir/bin/coordinator . +COPY --from=builder /build-dir/bin/logservice . +ENV PATH=$PATH:./ + +CMD /bin/bash diff --git a/go/Dockerfile.migration b/go/Dockerfile.migration new file mode 100644 index 00000000000..eab472c7324 --- /dev/null +++ b/go/Dockerfile.migration @@ -0,0 +1,9 @@ +FROM debian:bookworm-slim + +RUN apt update +RUN apt upgrade -y +RUN apt install -y curl jq +RUN curl -sSf https://atlasgo.sh | sh -s -- --community + +COPY ./go/migrations migrations +COPY ./go/atlas.hcl atlas.hcl diff --git a/go/coordinator/Makefile b/go/Makefile similarity index 56% rename from go/coordinator/Makefile rename to go/Makefile index 8fb52e4bb74..8bd6606d50b 100644 --- a/go/coordinator/Makefile +++ b/go/Makefile @@ -1,9 +1,11 @@ .PHONY: build build: - go build -v -o bin/chroma ./cmd + go build -v -o bin/coordinator ./cmd/coordinator/ + go build -v -o bin/logservice ./cmd/logservice/ test: build - go test -cover -race ./... + go test -race -cover ./... + lint: #brew install golangci-lint diff --git a/go/README.md b/go/README.md new file mode 100644 index 00000000000..0cfe009a428 --- /dev/null +++ b/go/README.md @@ -0,0 +1,19 @@ +# Set up Local Postgres + +- Install Postgres on Mac + - `brew install postgresql@14` +- Start & Stop + - `brew services start postgresql` + - `brew services stop postgresql` +- create testing db + - terminal: `psql postgres` + - postgres=# `create role chroma with login password 'chroma';` + - postgres=# `alter role chroma with superuser;` + - postgres=# `create database chroma;` +- Set postgres ENV Vars + Several tests (such as record_log_service_test.go) require the following environment variables to be set: + - `export POSTGRES_HOST=localhost` + - `export POSTGRES_PORT=5432` +- Atlas schema migration + - [~/chroma/go]: `atlas migrate diff --env dev` + - [~/chroma/go]: `atlas --env dev migrate apply --url "postgres://chroma:chroma@localhost:5432/chroma?sslmode=disable"` diff --git a/go/coordinator/atlas.hcl b/go/atlas.hcl similarity index 73% rename from go/coordinator/atlas.hcl rename to go/atlas.hcl index 2883c58d65e..ea5aa2d0916 100644 --- a/go/coordinator/atlas.hcl +++ b/go/atlas.hcl @@ -5,14 +5,14 @@ data "external_schema" "gorm" { "-mod=mod", "ariga.io/atlas-provider-gorm", "load", - "--path", "./internal/metastore/db/dbmodel", + "--path", "./pkg/metastore/db/dbmodel", "--dialect", "postgres", ] } -env "gorm" { +env "dev" { src = data.external_schema.gorm.url - dev = "postgres://localhost:5432/dev?sslmode=disable" + dev = "postgres://localhost:5432/chroma?sslmode=disable" migration { dir = "file://migrations" } diff --git a/go/cmd/coordinator/cmd.go b/go/cmd/coordinator/cmd.go new file mode 100644 index 00000000000..24a1093aad8 --- /dev/null +++ b/go/cmd/coordinator/cmd.go @@ -0,0 +1,72 @@ +package main + +import ( + "io" + "time" + + "github.com/chroma-core/chroma/go/pkg/coordinator/grpc" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + + "github.com/chroma-core/chroma/go/cmd/flag" + "github.com/chroma-core/chroma/go/pkg/utils" + "github.com/spf13/cobra" +) + +var ( + conf = grpc.Config{ + GrpcConfig: &grpcutils.GrpcConfig{}, + } + + Cmd = &cobra.Command{ + Use: "coordinator", + Short: "Start a coordinator", + Long: `Long description`, + Run: exec, + } +) + +func init() { + + // GRPC + flag.GRPCAddr(Cmd, &conf.GrpcConfig.BindAddress) + + // System Catalog + Cmd.Flags().StringVar(&conf.SystemCatalogProvider, "system-catalog-provider", "database", "System catalog provider") + Cmd.Flags().StringVar(&conf.DBConfig.Username, "username", "chroma", "MetaTable username") + Cmd.Flags().StringVar(&conf.DBConfig.Password, "password", "chroma", "MetaTable password") + Cmd.Flags().StringVar(&conf.DBConfig.Address, "db-address", "postgres", "MetaTable db address") + Cmd.Flags().IntVar(&conf.DBConfig.Port, "db-port", 5432, "MetaTable db port") + Cmd.Flags().StringVar(&conf.DBConfig.DBName, "db-name", "chroma", "MetaTable db name") + Cmd.Flags().IntVar(&conf.DBConfig.MaxIdleConns, "max-idle-conns", 10, "MetaTable max idle connections") + Cmd.Flags().IntVar(&conf.DBConfig.MaxOpenConns, "max-open-conns", 10, "MetaTable max open connections") + Cmd.Flags().StringVar(&conf.DBConfig.SslMode, "ssl-mode", "disable", "SSL mode for database connection") + + // Pulsar + Cmd.Flags().StringVar(&conf.PulsarAdminURL, "pulsar-admin-url", "http://localhost:8080", "Pulsar admin url") + Cmd.Flags().StringVar(&conf.PulsarURL, "pulsar-url", "pulsar://localhost:6650", "Pulsar url") + Cmd.Flags().StringVar(&conf.PulsarTenant, "pulsar-tenant", "default", "Pulsar tenant") + Cmd.Flags().StringVar(&conf.PulsarNamespace, "pulsar-namespace", "default", "Pulsar namespace") + + // Notification + Cmd.Flags().StringVar(&conf.NotificationStoreProvider, "notification-store-provider", "memory", "Notification store provider") + Cmd.Flags().StringVar(&conf.NotifierProvider, "notifier-provider", "memory", "Notifier provider") + Cmd.Flags().StringVar(&conf.NotificationTopic, "notification-topic", "chroma-notification", "Notification topic") + + // Memberlist + Cmd.Flags().StringVar(&conf.KubernetesNamespace, "kubernetes-namespace", "chroma", "Kubernetes namespace") + + // Query service memberlist + Cmd.Flags().StringVar(&conf.QueryServiceMemberlistName, "query-memberlist-name", "query-service-memberlist", "Query service memberlist name") + Cmd.Flags().StringVar(&conf.QueryServicePodLabel, "query-pod-label", "query-service", "Query pod label") + Cmd.Flags().DurationVar(&conf.WatchInterval, "watch-interval", 60*time.Second, "Watch interval") + + // Compaction service Memberlist + Cmd.Flags().StringVar(&conf.CompactionServiceMemberlistName, "compaction-memberlist-name", "compaction-service-memberlist", "Compaction memberlist name") + Cmd.Flags().StringVar(&conf.CompactionServicePodLabel, "compaction-pod-label", "compaction-service", "Compaction pod label") +} + +func exec(*cobra.Command, []string) { + utils.RunProcess(func() (io.Closer, error) { + return grpc.New(conf) + }) +} diff --git a/go/coordinator/cmd/main.go b/go/cmd/coordinator/main.go similarity index 77% rename from go/coordinator/cmd/main.go rename to go/cmd/coordinator/main.go index 0b7cfa7b54d..53fa1bd0967 100644 --- a/go/coordinator/cmd/main.go +++ b/go/cmd/coordinator/main.go @@ -4,8 +4,7 @@ import ( "fmt" "os" - "github.com/chroma/chroma-coordinator/cmd/grpccoordinator" - "github.com/chroma/chroma-coordinator/internal/utils" + "github.com/chroma-core/chroma/go/pkg/utils" "github.com/rs/zerolog" "github.com/spf13/cobra" "go.uber.org/automaxprocs/maxprocs" @@ -20,7 +19,7 @@ var ( ) func init() { - rootCmd.AddCommand(grpccoordinator.Cmd) + rootCmd.AddCommand(Cmd) } func main() { diff --git a/go/coordinator/cmd/flag/flag.go b/go/cmd/flag/flag.go similarity index 100% rename from go/coordinator/cmd/flag/flag.go rename to go/cmd/flag/flag.go diff --git a/go/cmd/logservice/cmd.go b/go/cmd/logservice/cmd.go new file mode 100644 index 00000000000..24d7adab5e5 --- /dev/null +++ b/go/cmd/logservice/cmd.go @@ -0,0 +1,46 @@ +package main + +import ( + "github.com/chroma-core/chroma/go/cmd/flag" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/chroma-core/chroma/go/pkg/logservice/grpc" + "github.com/chroma-core/chroma/go/pkg/utils" + "github.com/spf13/cobra" + "io" +) + +var ( + conf = grpc.Config{ + GrpcConfig: &grpcutils.GrpcConfig{}, + } + + Cmd = &cobra.Command{ + Use: "logservice", + Short: "Start a logservice service", + Long: `RecordLog root command`, + Run: exec, + } +) + +func init() { + // GRPC + flag.GRPCAddr(Cmd, &conf.GrpcConfig.BindAddress) + Cmd.Flags().BoolVar(&conf.StartGrpc, "start-grpc", true, "start grpc server or not") + + // DB provider + Cmd.Flags().StringVar(&conf.DBProvider, "db-provider", "postgres", "DB provider") + + // DB dev + Cmd.Flags().StringVar(&conf.DBConfig.Address, "db-host", "postgres", "DB host") + Cmd.Flags().IntVar(&conf.DBConfig.Port, "db-port", 5432, "DB port") + Cmd.Flags().StringVar(&conf.DBConfig.Username, "db-user", "chroma", "DB user") + Cmd.Flags().StringVar(&conf.DBConfig.Password, "db-password", "chroma", "DB password") + Cmd.Flags().StringVar(&conf.DBConfig.DBName, "db-name", "chroma", "DB name") + Cmd.Flags().StringVar(&conf.DBConfig.SslMode, "ssl-mode", "disable", "SSL mode for database connection") +} + +func exec(*cobra.Command, []string) { + utils.RunProcess(func() (io.Closer, error) { + return grpc.New(conf) + }) +} diff --git a/go/cmd/logservice/main.go b/go/cmd/logservice/main.go new file mode 100644 index 00000000000..5a9e8cb7def --- /dev/null +++ b/go/cmd/logservice/main.go @@ -0,0 +1,36 @@ +package main + +import ( + "fmt" + "os" + + "github.com/chroma-core/chroma/go/pkg/utils" + "github.com/rs/zerolog" + "github.com/spf13/cobra" + "go.uber.org/automaxprocs/maxprocs" +) + +var ( + rootCmd = &cobra.Command{ + Use: "logservice", + Short: "RecordLog root command", + Long: `RecordLog root command`, + } +) + +func init() { + rootCmd.AddCommand(Cmd) +} + +func main() { + utils.LogLevel = zerolog.DebugLevel + utils.ConfigureLogger() + if _, err := maxprocs.Set(); err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + if err := rootCmd.Execute(); err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} diff --git a/go/coordinator/Dockerfile b/go/coordinator/Dockerfile deleted file mode 100644 index a86f5cc258f..00000000000 --- a/go/coordinator/Dockerfile +++ /dev/null @@ -1,31 +0,0 @@ -FROM golang:1.20-alpine3.18 as build - -RUN apk add --no-cache make git build-base bash - -ENV PATH=$PATH:/go/bin -ADD ./go/coordinator /src/chroma-coordinator - -RUN cd /src/chroma-coordinator \ - && make - -FROM alpine:3.17.3 - -RUN apk add --no-cache bash bash-completion jq findutils - -# As of 6 Dec 2023, the atlas package isn't in Alpine's main package manager, only -# testing. So we have to add the testing repository to get it. -RUN apk add \ - --no-cache \ - --repository http://dl-cdn.alpinelinux.org/alpine/edge/testing \ - --repository http://dl-cdn.alpinelinux.org/alpine/edge/main \ - atlas - -RUN mkdir /chroma-coordinator -WORKDIR /chroma-coordinator - -COPY --from=build /src/chroma-coordinator/bin/chroma /chroma-coordinator/bin/chroma -ENV PATH=$PATH:/chroma-coordinator/bin - -COPY --from=build /src/chroma-coordinator/migrations /chroma-coordinator/migrations - -CMD /bin/bash diff --git a/go/coordinator/cmd/grpccoordinator/cmd.go b/go/coordinator/cmd/grpccoordinator/cmd.go deleted file mode 100644 index 8859790b56c..00000000000 --- a/go/coordinator/cmd/grpccoordinator/cmd.go +++ /dev/null @@ -1,64 +0,0 @@ -package grpccoordinator - -import ( - "io" - "time" - - "github.com/chroma/chroma-coordinator/cmd/flag" - "github.com/chroma/chroma-coordinator/internal/grpccoordinator" - "github.com/chroma/chroma-coordinator/internal/grpccoordinator/grpcutils" - "github.com/chroma/chroma-coordinator/internal/utils" - "github.com/spf13/cobra" -) - -var ( - conf = grpccoordinator.Config{ - GrpcConfig: &grpcutils.GrpcConfig{}, - } - - Cmd = &cobra.Command{ - Use: "coordinator", - Short: "Start a coordinator", - Long: `Long description`, - Run: exec, - } -) - -func init() { - - // GRPC - flag.GRPCAddr(Cmd, &conf.GrpcConfig.BindAddress) - - // System Catalog - Cmd.Flags().StringVar(&conf.SystemCatalogProvider, "system-catalog-provider", "memory", "System catalog provider") - Cmd.Flags().StringVar(&conf.Username, "username", "root", "MetaTable username") - Cmd.Flags().StringVar(&conf.Password, "password", "", "MetaTable password") - Cmd.Flags().StringVar(&conf.Address, "db-address", "127.0.0.1", "MetaTable db address") - Cmd.Flags().IntVar(&conf.Port, "db-port", 5432, "MetaTable db port") - Cmd.Flags().StringVar(&conf.DBName, "db-name", "", "MetaTable db name") - Cmd.Flags().IntVar(&conf.MaxIdleConns, "max-idle-conns", 10, "MetaTable max idle connections") - Cmd.Flags().IntVar(&conf.MaxOpenConns, "max-open-conns", 10, "MetaTable max open connections") - - // Pulsar - Cmd.Flags().StringVar(&conf.PulsarAdminURL, "pulsar-admin-url", "http://localhost:8080", "Pulsar admin url") - Cmd.Flags().StringVar(&conf.PulsarURL, "pulsar-url", "pulsar://localhost:6650", "Pulsar url") - Cmd.Flags().StringVar(&conf.PulsarTenant, "pulsar-tenant", "default", "Pulsar tenant") - Cmd.Flags().StringVar(&conf.PulsarNamespace, "pulsar-namespace", "default", "Pulsar namespace") - - // Notification - Cmd.Flags().StringVar(&conf.NotificationStoreProvider, "notification-store-provider", "memory", "Notification store provider") - Cmd.Flags().StringVar(&conf.NotifierProvider, "notifier-provider", "memory", "Notifier provider") - Cmd.Flags().StringVar(&conf.NotificationTopic, "notification-topic", "chroma-notification", "Notification topic") - - // Memberlist - Cmd.Flags().StringVar(&conf.KubernetesNamespace, "kubernetes-namespace", "chroma", "Kubernetes namespace") - Cmd.Flags().StringVar(&conf.WorkerMemberlistName, "worker-memberlist-name", "worker-memberlist", "Worker memberlist name") - Cmd.Flags().StringVar(&conf.AssignmentPolicy, "assignment-policy", "rendezvous", "Assignment policy") - Cmd.Flags().DurationVar(&conf.WatchInterval, "watch-interval", 60*time.Second, "Watch interval") -} - -func exec(*cobra.Command, []string) { - utils.RunProcess(func() (io.Closer, error) { - return grpccoordinator.New(conf) - }) -} diff --git a/go/coordinator/internal/coordinator/apis_test.go b/go/coordinator/internal/coordinator/apis_test.go deleted file mode 100644 index 62ff01ecec0..00000000000 --- a/go/coordinator/internal/coordinator/apis_test.go +++ /dev/null @@ -1,957 +0,0 @@ -package coordinator - -import ( - "context" - "sort" - "testing" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbcore" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/google/uuid" - "github.com/stretchr/testify/assert" - "pgregory.net/rapid" -) - -// TODO: This is not complete yet. We need to add more tests for the other APIs. -// We will deprecate the example based tests once we have enough tests here. -func testCollection(t *rapid.T) { - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewSimpleAssignmentPolicy("test-tenant", "test-topic") - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - t.Repeat(map[string]func(*rapid.T){ - "create_collection": func(t *rapid.T) { - stringValue := generateCollectionStringMetadataValue(t) - intValue := generateCollectionInt64MetadataValue(t) - floatValue := generateCollectionFloat64MetadataValue(t) - - metadata := model.NewCollectionMetadata[model.CollectionMetadataValueType]() - metadata.Add("string_value", stringValue) - metadata.Add("int_value", intValue) - metadata.Add("float_value", floatValue) - - collection := rapid.Custom[*model.CreateCollection](func(t *rapid.T) *model.CreateCollection { - return &model.CreateCollection{ - ID: types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "collection_id")), - Name: rapid.String().Draw(t, "collection_name"), - Metadata: nil, - } - }).Draw(t, "collection") - - _, err := c.CreateCollection(ctx, collection) - if err != nil { - if err == common.ErrCollectionNameEmpty && collection.Name == "" { - t.Logf("expected error for empty collection name") - } else if err == common.ErrCollectionTopicEmpty { - t.Logf("expected error for empty collection topic") - } else { - t.Fatalf("error creating collection: %v", err) - } - } - if err == nil { - // verify the correctness - collectionList, err := c.GetCollections(ctx, collection.ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - if err != nil { - t.Fatalf("error getting collections: %v", err) - } - if len(collectionList) != 1 { - t.Fatalf("More than 1 collection with the same collection id") - } - for _, collectionReturned := range collectionList { - if collection.ID != collectionReturned.ID { - t.Fatalf("collection id is the right value") - } - } - } - }, - }) -} - -func testSegment(t *rapid.T) { - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewSimpleAssignmentPolicy("test-tenant", "test-topic") - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - - stringValue := generateSegmentStringMetadataValue(t) - intValue := generateSegmentInt64MetadataValue(t) - floatValue := generateSegmentFloat64MetadataValue(t) - - metadata := model.NewSegmentMetadata[model.SegmentMetadataValueType]() - metadata.Set("string_value", stringValue) - metadata.Set("int_value", intValue) - metadata.Set("float_value", floatValue) - - testTopic := "test-segment-topic" - t.Repeat(map[string]func(*rapid.T){ - "create_segment": func(t *rapid.T) { - segment := rapid.Custom[*model.CreateSegment](func(t *rapid.T) *model.CreateSegment { - return &model.CreateSegment{ - ID: types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "segment_id")), - Type: "test-segment-type", - Scope: "test-segment-scope", - Topic: &testTopic, - Metadata: nil, - CollectionID: types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "collection_id")), - } - }).Draw(t, "segment") - - err := c.CreateSegment(ctx, segment) - if err != nil { - t.Fatalf("error creating segment: %v", err) - } - }, - }) -} - -func generateCollectionStringMetadataValue(t *rapid.T) model.CollectionMetadataValueType { - return &model.CollectionMetadataValueStringType{ - Value: rapid.String().Draw(t, "string_value"), - } -} - -func generateCollectionInt64MetadataValue(t *rapid.T) model.CollectionMetadataValueType { - return &model.CollectionMetadataValueInt64Type{ - Value: rapid.Int64().Draw(t, "int_value"), - } -} - -func generateCollectionFloat64MetadataValue(t *rapid.T) model.CollectionMetadataValueType { - return &model.CollectionMetadataValueFloat64Type{ - Value: rapid.Float64().Draw(t, "float_value"), - } -} - -func generateSegmentStringMetadataValue(t *rapid.T) model.SegmentMetadataValueType { - return &model.SegmentMetadataValueStringType{ - Value: rapid.String().Draw(t, "string_value"), - } -} - -func generateSegmentInt64MetadataValue(t *rapid.T) model.SegmentMetadataValueType { - return &model.SegmentMetadataValueInt64Type{ - Value: rapid.Int64().Draw(t, "int_value"), - } -} - -func generateSegmentFloat64MetadataValue(t *rapid.T) model.SegmentMetadataValueType { - return &model.SegmentMetadataValueFloat64Type{ - Value: rapid.Float64().Draw(t, "float_value"), - } -} - -func TestAPIs(t *testing.T) { - // rapid.Check(t, testCollection) - // rapid.Check(t, testSegment) -} - -func SampleCollections(t *testing.T, tenantID string, databaseName string) []*model.Collection { - dimension := int32(128) - metadata1 := model.NewCollectionMetadata[model.CollectionMetadataValueType]() - metadata1.Add("test_str", &model.CollectionMetadataValueStringType{Value: "str1"}) - metadata1.Add("test_int", &model.CollectionMetadataValueInt64Type{Value: 1}) - metadata1.Add("test_float", &model.CollectionMetadataValueFloat64Type{Value: 1.3}) - - metadata2 := model.NewCollectionMetadata[model.CollectionMetadataValueType]() - metadata2.Add("test_str", &model.CollectionMetadataValueStringType{Value: "str2"}) - metadata2.Add("test_int", &model.CollectionMetadataValueInt64Type{Value: 2}) - metadata2.Add("test_float", &model.CollectionMetadataValueFloat64Type{Value: 2.3}) - - metadata3 := model.NewCollectionMetadata[model.CollectionMetadataValueType]() - metadata3.Add("test_str", &model.CollectionMetadataValueStringType{Value: "str3"}) - metadata3.Add("test_int", &model.CollectionMetadataValueInt64Type{Value: 3}) - metadata3.Add("test_float", &model.CollectionMetadataValueFloat64Type{Value: 3.3}) - sampleCollections := []*model.Collection{ - { - ID: types.MustParse("93ffe3ec-0107-48d4-8695-51f978c509dc"), - Name: "test_collection_1", - Topic: "test_topic_1", - Metadata: metadata1, - Dimension: &dimension, - Created: true, - TenantID: tenantID, - DatabaseName: databaseName, - }, - { - ID: types.MustParse("f444f1d7-d06c-4357-ac22-5a4a1f92d761"), - Name: "test_collection_2", - Topic: "test_topic_2", - Metadata: metadata2, - Dimension: nil, - Created: true, - TenantID: tenantID, - DatabaseName: databaseName, - }, - { - ID: types.MustParse("43babc1a-e403-4a50-91a9-16621ba29ab0"), - Name: "test_collection_3", - Topic: "test_topic_3", - Metadata: metadata3, - Dimension: nil, - Created: true, - TenantID: tenantID, - DatabaseName: databaseName, - }, - } - return sampleCollections -} - -type MockAssignmentPolicy struct { - collections []*model.Collection -} - -func NewMockAssignmentPolicy(collecions []*model.Collection) *MockAssignmentPolicy { - return &MockAssignmentPolicy{ - collections: collecions, - } -} - -func (m *MockAssignmentPolicy) AssignCollection(collectionID types.UniqueID) (string, error) { - for _, collection := range m.collections { - if collection.ID == collectionID { - return collection.Topic, nil - } - } - return "", common.ErrCollectionNotFound -} - -func TestCreateGetDeleteCollections(t *testing.T) { - - sampleCollections := SampleCollections(t, common.DefaultTenant, common.DefaultDatabase) - - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - - for _, collection := range sampleCollections { - c.CreateCollection(ctx, &model.CreateCollection{ - ID: collection.ID, - Name: collection.Name, - Topic: collection.Topic, - Metadata: collection.Metadata, - Dimension: collection.Dimension, - TenantID: collection.TenantID, - DatabaseName: collection.DatabaseName, - }) - } - - results, err := c.GetCollections(ctx, types.NilUniqueID(), nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - - sort.Slice(results, func(i, j int) bool { - return results[i].Name < results[j].Name - }) - - assert.Equal(t, sampleCollections, results) - - // Duplicate create fails - _, err = c.CreateCollection(ctx, &model.CreateCollection{ - ID: sampleCollections[0].ID, - Name: sampleCollections[0].Name, - TenantID: common.DefaultTenant, - DatabaseName: common.DefaultDatabase, - }) - assert.Error(t, err) - - // Find by name - for _, collection := range sampleCollections { - result, err := c.GetCollections(ctx, types.NilUniqueID(), &collection.Name, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{collection}, result) - } - - // Find by topic - for _, collection := range sampleCollections { - result, err := c.GetCollections(ctx, types.NilUniqueID(), nil, &collection.Topic, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{collection}, result) - } - - // Find by id - for _, collection := range sampleCollections { - result, err := c.GetCollections(ctx, collection.ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{collection}, result) - } - - // Find by id and topic (positive case) - for _, collection := range sampleCollections { - result, err := c.GetCollections(ctx, collection.ID, nil, &collection.Topic, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{collection}, result) - } - - // find by id and topic (negative case) - for _, collection := range sampleCollections { - otherTopic := "other topic" - result, err := c.GetCollections(ctx, collection.ID, nil, &otherTopic, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Empty(t, result) - } - - // Delete - c1 := sampleCollections[0] - deleteCollection := &model.DeleteCollection{ - ID: c1.ID, - DatabaseName: common.DefaultDatabase, - TenantID: common.DefaultTenant, - } - err = c.DeleteCollection(ctx, deleteCollection) - assert.NoError(t, err) - - results, err = c.GetCollections(ctx, types.NilUniqueID(), nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - - assert.NotContains(t, results, c1) - assert.Len(t, results, len(sampleCollections)-1) - assert.ElementsMatch(t, results, sampleCollections[1:]) - byIDResult, err := c.GetCollections(ctx, c1.ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Empty(t, byIDResult) - - // Duplicate delete throws an exception - err = c.DeleteCollection(ctx, deleteCollection) - assert.Error(t, err) -} - -func TestUpdateCollections(t *testing.T) { - sampleCollections := SampleCollections(t, common.DefaultTenant, common.DefaultDatabase) - - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - - coll := &model.Collection{ - Name: sampleCollections[0].Name, - ID: sampleCollections[0].ID, - Topic: sampleCollections[0].Topic, - Metadata: sampleCollections[0].Metadata, - Dimension: sampleCollections[0].Dimension, - Created: false, - TenantID: sampleCollections[0].TenantID, - DatabaseName: sampleCollections[0].DatabaseName, - } - - c.CreateCollection(ctx, &model.CreateCollection{ - ID: coll.ID, - Name: coll.Name, - Topic: coll.Topic, - Metadata: coll.Metadata, - Dimension: coll.Dimension, - TenantID: coll.TenantID, - DatabaseName: coll.DatabaseName, - }) - - // Update name - coll.Name = "new_name" - result, err := c.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Name: &coll.Name}) - assert.NoError(t, err) - assert.Equal(t, coll, result) - resultList, err := c.GetCollections(ctx, types.NilUniqueID(), &coll.Name, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{coll}, resultList) - - // Update topic - coll.Topic = "new_topic" - result, err = c.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Topic: &coll.Topic}) - assert.NoError(t, err) - assert.Equal(t, coll, result) - resultList, err = c.GetCollections(ctx, types.NilUniqueID(), nil, &coll.Topic, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{coll}, resultList) - - // Update dimension - newDimension := int32(128) - coll.Dimension = &newDimension - result, err = c.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Dimension: coll.Dimension}) - assert.NoError(t, err) - assert.Equal(t, coll, result) - resultList, err = c.GetCollections(ctx, coll.ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{coll}, resultList) - - // Reset the metadata - newMetadata := model.NewCollectionMetadata[model.CollectionMetadataValueType]() - newMetadata.Add("test_str2", &model.CollectionMetadataValueStringType{Value: "str2"}) - coll.Metadata = newMetadata - result, err = c.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Metadata: coll.Metadata}) - assert.NoError(t, err) - assert.Equal(t, coll, result) - resultList, err = c.GetCollections(ctx, coll.ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{coll}, resultList) - - // Delete all metadata keys - coll.Metadata = nil - result, err = c.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Metadata: coll.Metadata, ResetMetadata: true}) - assert.NoError(t, err) - assert.Equal(t, coll, result) - resultList, err = c.GetCollections(ctx, coll.ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, []*model.Collection{coll}, resultList) -} - -func TestCreateUpdateWithDatabase(t *testing.T) { - sampleCollections := SampleCollections(t, common.DefaultTenant, common.DefaultDatabase) - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - _, err = c.CreateDatabase(ctx, &model.CreateDatabase{ - ID: types.MustParse("00000000-d7d7-413b-92e1-731098a6e492").String(), - Name: "new_database", - Tenant: common.DefaultTenant, - }) - assert.NoError(t, err) - - c.CreateCollection(ctx, &model.CreateCollection{ - ID: sampleCollections[0].ID, - Name: sampleCollections[0].Name, - Topic: sampleCollections[0].Topic, - Metadata: sampleCollections[0].Metadata, - Dimension: sampleCollections[0].Dimension, - TenantID: sampleCollections[0].TenantID, - DatabaseName: "new_database", - }) - - c.CreateCollection(ctx, &model.CreateCollection{ - ID: sampleCollections[1].ID, - Name: sampleCollections[1].Name, - Topic: sampleCollections[1].Topic, - Metadata: sampleCollections[1].Metadata, - Dimension: sampleCollections[1].Dimension, - TenantID: sampleCollections[1].TenantID, - DatabaseName: sampleCollections[1].DatabaseName, - }) - - newName1 := "new_name_1" - c.UpdateCollection(ctx, &model.UpdateCollection{ - ID: sampleCollections[1].ID, - Name: &newName1, - }) - - result, err := c.GetCollections(ctx, sampleCollections[1].ID, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, 1, len(result)) - assert.Equal(t, "new_name_1", result[0].Name) - - newName0 := "new_name_0" - c.UpdateCollection(ctx, &model.UpdateCollection{ - ID: sampleCollections[0].ID, - Name: &newName0, - }) - result, err = c.GetCollections(ctx, sampleCollections[0].ID, nil, nil, common.DefaultTenant, "new_database") - assert.NoError(t, err) - assert.Equal(t, 1, len(result)) - assert.Equal(t, "new_name_0", result[0].Name) -} - -func TestGetMultipleWithDatabase(t *testing.T) { - sampleCollections := SampleCollections(t, common.DefaultTenant, "new_database") - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - _, err = c.CreateDatabase(ctx, &model.CreateDatabase{ - ID: types.MustParse("00000000-d7d7-413b-92e1-731098a6e492").String(), - Name: "new_database", - Tenant: common.DefaultTenant, - }) - assert.NoError(t, err) - - for _, collection := range sampleCollections { - c.CreateCollection(ctx, &model.CreateCollection{ - ID: collection.ID, - Name: collection.Name, - Topic: collection.Topic, - Metadata: collection.Metadata, - Dimension: collection.Dimension, - TenantID: common.DefaultTenant, - DatabaseName: "new_database", - }) - } - result, err := c.GetCollections(ctx, types.NilUniqueID(), nil, nil, common.DefaultTenant, "new_database") - assert.NoError(t, err) - assert.Equal(t, len(sampleCollections), len(result)) - sort.Slice(result, func(i, j int) bool { - return result[i].Name < result[j].Name - }) - assert.Equal(t, sampleCollections, result) - - result, err = c.GetCollections(ctx, types.NilUniqueID(), nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Equal(t, 0, len(result)) -} - -func TestCreateDatabaseWithTenants(t *testing.T) { - sampleCollections := SampleCollections(t, common.DefaultTenant, common.DefaultDatabase) - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - - // Create a new tenant - _, err = c.CreateTenant(ctx, &model.CreateTenant{ - Name: "tenant1", - }) - assert.NoError(t, err) - - // Create tenant that already exits and expect an error - _, err = c.CreateTenant(ctx, &model.CreateTenant{ - Name: "tenant1", - }) - assert.Error(t, err) - - // Create tenant that already exits and expect an error - _, err = c.CreateTenant(ctx, &model.CreateTenant{ - Name: common.DefaultTenant, - }) - assert.Error(t, err) - - // Create a new database within this tenant and also in the default tenant - _, err = c.CreateDatabase(ctx, &model.CreateDatabase{ - ID: types.MustParse("33333333-d7d7-413b-92e1-731098a6e492").String(), - Name: "new_database", - Tenant: "tenant1", - }) - assert.NoError(t, err) - - _, err = c.CreateDatabase(ctx, &model.CreateDatabase{ - ID: types.MustParse("44444444-d7d7-413b-92e1-731098a6e492").String(), - Name: "new_database", - Tenant: common.DefaultTenant, - }) - assert.NoError(t, err) - - // Create a new collection in the new tenant - _, err = c.CreateCollection(ctx, &model.CreateCollection{ - ID: sampleCollections[0].ID, - Name: sampleCollections[0].Name, - Topic: sampleCollections[0].Topic, - Metadata: sampleCollections[0].Metadata, - Dimension: sampleCollections[0].Dimension, - TenantID: "tenant1", - DatabaseName: "new_database", - }) - assert.NoError(t, err) - - // Create a new collection in the default tenant - c.CreateCollection(ctx, &model.CreateCollection{ - ID: sampleCollections[1].ID, - Name: sampleCollections[1].Name, - Topic: sampleCollections[1].Topic, - Metadata: sampleCollections[1].Metadata, - Dimension: sampleCollections[1].Dimension, - TenantID: common.DefaultTenant, - DatabaseName: "new_database", - }) - - // Check that both tenants have the correct collections - expected := []*model.Collection{sampleCollections[0]} - expected[0].TenantID = "tenant1" - expected[0].DatabaseName = "new_database" - result, err := c.GetCollections(ctx, types.NilUniqueID(), nil, nil, "tenant1", "new_database") - assert.NoError(t, err) - assert.Equal(t, 1, len(result)) - assert.Equal(t, expected[0], result[0]) - - expected = []*model.Collection{sampleCollections[1]} - expected[0].TenantID = common.DefaultTenant - expected[0].DatabaseName = "new_database" - result, err = c.GetCollections(ctx, types.NilUniqueID(), nil, nil, common.DefaultTenant, "new_database") - assert.NoError(t, err) - assert.Equal(t, 1, len(result)) - assert.Equal(t, expected[0], result[0]) - - // A new tenant DOES NOT have a default database. This does not error, instead 0 - // results are returned - result, err = c.GetCollections(ctx, types.NilUniqueID(), nil, nil, "tenant1", common.DefaultDatabase) - assert.Error(t, err) - assert.Nil(t, result) -} - -func TestCreateGetDeleteTenants(t *testing.T) { - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(nil) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - - // Create a new tenant - _, err = c.CreateTenant(ctx, &model.CreateTenant{ - Name: "tenant1", - }) - assert.NoError(t, err) - - // Create tenant that already exits and expect an error - _, err = c.CreateTenant(ctx, &model.CreateTenant{ - Name: "tenant1", - }) - assert.Error(t, err) - - // Create tenant that already exits and expect an error - _, err = c.CreateTenant(ctx, &model.CreateTenant{ - Name: common.DefaultTenant, - }) - assert.Error(t, err) - - // Get the tenant and check that it exists - result, err := c.GetTenant(ctx, &model.GetTenant{Name: "tenant1"}) - assert.NoError(t, err) - assert.Equal(t, "tenant1", result.Name) - - // Get a tenant that does not exist and expect an error - _, err = c.GetTenant(ctx, &model.GetTenant{Name: "tenant2"}) - assert.Error(t, err) - - // Create a new database within this tenant - _, err = c.CreateDatabase(ctx, &model.CreateDatabase{ - ID: types.MustParse("33333333-d7d7-413b-92e1-731098a6e492").String(), - Name: "new_database", - Tenant: "tenant1", - }) - assert.NoError(t, err) - - // Get the database and check that it exists - databaseResult, err := c.GetDatabase(ctx, &model.GetDatabase{ - Name: "new_database", - Tenant: "tenant1", - }) - assert.NoError(t, err) - assert.Equal(t, "new_database", databaseResult.Name) - assert.Equal(t, "tenant1", databaseResult.Tenant) - - // Get a database that does not exist in a tenant that does exist and expect an error - _, err = c.GetDatabase(ctx, &model.GetDatabase{ - Name: "new_database1", - Tenant: "tenant1", - }) - assert.Error(t, err) - - // Get a database that does not exist in a tenant that does not exist and expect an - // error - _, err = c.GetDatabase(ctx, &model.GetDatabase{ - Name: "new_database1", - Tenant: "tenant2", - }) - assert.Error(t, err) -} - -func SampleSegments(t *testing.T, sampleCollections []*model.Collection) []*model.Segment { - metadata1 := model.NewSegmentMetadata[model.SegmentMetadataValueType]() - metadata1.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str1"}) - metadata1.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 1}) - metadata1.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 1.3}) - - metadata2 := model.NewSegmentMetadata[model.SegmentMetadataValueType]() - metadata2.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str2"}) - metadata2.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 2}) - metadata2.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 2.3}) - - metadata3 := model.NewSegmentMetadata[model.SegmentMetadataValueType]() - metadata3.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str3"}) - metadata3.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 3}) - metadata3.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 3.3}) - - testTopic2 := "test_topic_2" - testTopic3 := "test_topic_3" - sampleSegments := []*model.Segment{ - { - ID: types.MustParse("00000000-d7d7-413b-92e1-731098a6e492"), - Type: "test_type_a", - Topic: nil, - Scope: "VECTOR", - CollectionID: sampleCollections[0].ID, - Metadata: metadata1, - }, - { - ID: types.MustParse("11111111-d7d7-413b-92e1-731098a6e492"), - Type: "test_type_b", - Topic: &testTopic2, - Scope: "VECTOR", - CollectionID: sampleCollections[1].ID, - Metadata: metadata2, - }, - { - ID: types.MustParse("22222222-d7d7-413b-92e1-731098a6e492"), - Type: "test_type_b", - Topic: &testTopic3, - Scope: "METADATA", - CollectionID: types.NilUniqueID(), - Metadata: metadata3, // This segment is not assigned to any collection - }, - } - return sampleSegments -} - -func TestCreateGetDeleteSegments(t *testing.T) { - sampleCollections := SampleCollections(t, common.DefaultTenant, common.DefaultDatabase) - - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - - for _, collection := range sampleCollections { - c.CreateCollection(ctx, &model.CreateCollection{ - ID: collection.ID, - Name: collection.Name, - Topic: collection.Topic, - Metadata: collection.Metadata, - Dimension: collection.Dimension, - TenantID: collection.TenantID, - DatabaseName: collection.DatabaseName, - }) - } - - sampleSegments := SampleSegments(t, sampleCollections) - for _, segment := range sampleSegments { - c.CreateSegment(ctx, &model.CreateSegment{ - ID: segment.ID, - Type: segment.Type, - Topic: segment.Topic, - Scope: segment.Scope, - CollectionID: segment.CollectionID, - Metadata: segment.Metadata, - }) - } - - results, err := c.GetSegments(ctx, types.NilUniqueID(), nil, nil, nil, types.NilUniqueID()) - sort.Slice(results, func(i, j int) bool { - return results[i].ID.String() < results[j].ID.String() - }) - assert.NoError(t, err) - assert.Equal(t, sampleSegments, results) - - // Duplicate create fails - err = c.CreateSegment(ctx, &model.CreateSegment{ - ID: sampleSegments[0].ID, - Type: sampleSegments[0].Type, - Topic: sampleSegments[0].Topic, - Scope: sampleSegments[0].Scope, - CollectionID: sampleSegments[0].CollectionID, - Metadata: sampleSegments[0].Metadata, - }) - assert.Error(t, err) - - // Find by id - for _, segment := range sampleSegments { - result, err := c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - } - - // Find by type - testTypeA := "test_type_a" - result, err := c.GetSegments(ctx, types.NilUniqueID(), &testTypeA, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, sampleSegments[:1], result) - - testTypeB := "test_type_b" - result, err = c.GetSegments(ctx, types.NilUniqueID(), &testTypeB, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.ElementsMatch(t, result, sampleSegments[1:]) - - // Find by collection ID - result, err = c.GetSegments(ctx, types.NilUniqueID(), nil, nil, nil, sampleCollections[0].ID) - assert.NoError(t, err) - assert.Equal(t, sampleSegments[:1], result) - - // Find by type and collection ID (positive case) - result, err = c.GetSegments(ctx, types.NilUniqueID(), &testTypeA, nil, nil, sampleCollections[0].ID) - assert.NoError(t, err) - assert.Equal(t, sampleSegments[:1], result) - - // Find by type and collection ID (negative case) - result, err = c.GetSegments(ctx, types.NilUniqueID(), &testTypeB, nil, nil, sampleCollections[0].ID) - assert.NoError(t, err) - assert.Empty(t, result) - - // Delete - s1 := sampleSegments[0] - err = c.DeleteSegment(ctx, s1.ID) - assert.NoError(t, err) - - results, err = c.GetSegments(ctx, types.NilUniqueID(), nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.NotContains(t, results, s1) - assert.Len(t, results, len(sampleSegments)-1) - assert.ElementsMatch(t, results, sampleSegments[1:]) - - // Duplicate delete throws an exception - err = c.DeleteSegment(ctx, s1.ID) - assert.Error(t, err) -} - -func TestUpdateSegment(t *testing.T) { - sampleCollections := SampleCollections(t, common.DefaultTenant, common.DefaultDatabase) - - db := dbcore.ConfigDatabaseForTesting() - ctx := context.Background() - assignmentPolicy := NewMockAssignmentPolicy(sampleCollections) - c, err := NewCoordinator(ctx, assignmentPolicy, db, nil, nil) - if err != nil { - t.Fatalf("error creating coordinator: %v", err) - } - c.ResetState(ctx) - - testTopic := "test_topic_a" - - metadata := model.NewSegmentMetadata[model.SegmentMetadataValueType]() - metadata.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str1"}) - metadata.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 1}) - metadata.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 1.3}) - - segment := &model.Segment{ - ID: types.UniqueID(uuid.New()), - Type: "test_type_a", - Scope: "VECTOR", - Topic: &testTopic, - CollectionID: sampleCollections[0].ID, - Metadata: metadata, - } - - for _, collection := range sampleCollections { - _, err := c.CreateCollection(ctx, &model.CreateCollection{ - ID: collection.ID, - Name: collection.Name, - Topic: collection.Topic, - Metadata: collection.Metadata, - Dimension: collection.Dimension, - TenantID: collection.TenantID, - DatabaseName: collection.DatabaseName, - }) - - assert.NoError(t, err) - } - - c.CreateSegment(ctx, &model.CreateSegment{ - ID: segment.ID, - Type: segment.Type, - Topic: segment.Topic, - Scope: segment.Scope, - CollectionID: segment.CollectionID, - Metadata: segment.Metadata, - }) - - // Update topic to new value - newTopic := "new_topic" - segment.Topic = &newTopic - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Topic: segment.Topic, - }) - result, err := c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Update topic to None - segment.Topic = nil - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Topic: segment.Topic, - ResetTopic: true, - }) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Update collection to new value - segment.CollectionID = sampleCollections[1].ID - newCollecionID := segment.CollectionID.String() - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Collection: &newCollecionID, - }) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Update collection to None - segment.CollectionID = types.NilUniqueID() - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Collection: nil, - ResetCollection: true, - }) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Add a new metadata key - segment.Metadata.Set("test_str2", &model.SegmentMetadataValueStringType{Value: "str2"}) - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Metadata: segment.Metadata}) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Update a metadata key - segment.Metadata.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str3"}) - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Metadata: segment.Metadata}) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Delete a metadata key - segment.Metadata.Remove("test_str") - newMetadata := model.NewSegmentMetadata[model.SegmentMetadataValueType]() - newMetadata.Set("test_str", nil) - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Metadata: newMetadata}) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) - - // Delete all metadata keys - segment.Metadata = nil - c.UpdateSegment(ctx, &model.UpdateSegment{ - ID: segment.ID, - Metadata: segment.Metadata, - ResetMetadata: true}, - ) - result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Equal(t, []*model.Segment{segment}, result) -} diff --git a/go/coordinator/internal/coordinator/assignment_policy.go b/go/coordinator/internal/coordinator/assignment_policy.go deleted file mode 100644 index 6976d6a9652..00000000000 --- a/go/coordinator/internal/coordinator/assignment_policy.go +++ /dev/null @@ -1,77 +0,0 @@ -package coordinator - -import ( - "fmt" - - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/chroma/chroma-coordinator/internal/utils" -) - -type CollectionAssignmentPolicy interface { - AssignCollection(collectionID types.UniqueID) (string, error) -} - -// SimpleAssignmentPolicy is a simple assignment policy that assigns a 1 collection to 1 -// topic based on the id of the collection. -type SimpleAssignmentPolicy struct { - tenantID string - topicNS string -} - -func NewSimpleAssignmentPolicy(tenantID string, topicNS string) *SimpleAssignmentPolicy { - return &SimpleAssignmentPolicy{ - tenantID: tenantID, - topicNS: topicNS, - } -} - -func (s *SimpleAssignmentPolicy) AssignCollection(collectionID types.UniqueID) (string, error) { - return createTopicName(s.tenantID, s.topicNS, collectionID.String()), nil -} - -func createTopicName(tenantID string, topicNS string, log_name string) string { - return fmt.Sprintf("persistent://%s/%s/%s", tenantID, topicNS, log_name) -} - -// RendezvousAssignmentPolicy is an assignment policy that assigns a collection to a topic -// For now it assumes there are 16 topics and uses the rendezvous hashing algorithm to -// assign a collection to a topic. - -var Topics = [16]string{ - "chroma_log_0", - "chroma_log_1", - "chroma_log_2", - "chroma_log_3", - "chroma_log_4", - "chroma_log_5", - "chroma_log_6", - "chroma_log_7", - "chroma_log_8", - "chroma_log_9", - "chroma_log_10", - "chroma_log_11", - "chroma_log_12", - "chroma_log_13", - "chroma_log_14", - "chroma_log_15", -} - -type RendezvousAssignmentPolicy struct { - tenantID string - topicNS string -} - -func NewRendezvousAssignmentPolicy(tenantID string, topicNS string) *RendezvousAssignmentPolicy { - return &RendezvousAssignmentPolicy{ - tenantID: tenantID, - topicNS: topicNS, - } -} - -func (r *RendezvousAssignmentPolicy) AssignCollection(collectionID types.UniqueID) (string, error) { - assignment, error := utils.Assign(collectionID.String(), Topics[:], utils.Murmur3Hasher) - if error != nil { - return "", error - } - return createTopicName(r.tenantID, r.topicNS, assignment), nil -} diff --git a/go/coordinator/internal/coordinator/coordinator.go b/go/coordinator/internal/coordinator/coordinator.go deleted file mode 100644 index 2f3c02cff26..00000000000 --- a/go/coordinator/internal/coordinator/coordinator.go +++ /dev/null @@ -1,76 +0,0 @@ -package coordinator - -import ( - "context" - "log" - - "github.com/chroma/chroma-coordinator/internal/metastore" - "github.com/chroma/chroma-coordinator/internal/metastore/coordinator" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dao" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbcore" - "github.com/chroma/chroma-coordinator/internal/notification" - "github.com/chroma/chroma-coordinator/internal/types" - "gorm.io/gorm" -) - -var _ ICoordinator = (*Coordinator)(nil) - -// Coordinator is the implemenation of ICoordinator. It is the top level component. -// Currently, it only has the system catalog related APIs and will be extended to -// support other functionalities such as membership managed and propagation. -type Coordinator struct { - ctx context.Context - collectionAssignmentPolicy CollectionAssignmentPolicy - meta IMeta - notificationProcessor notification.NotificationProcessor -} - -func NewCoordinator(ctx context.Context, assignmentPolicy CollectionAssignmentPolicy, db *gorm.DB, notificationStore notification.NotificationStore, notifier notification.Notifier) (*Coordinator, error) { - s := &Coordinator{ - ctx: ctx, - collectionAssignmentPolicy: assignmentPolicy, - } - - notificationProcessor := notification.NewSimpleNotificationProcessor(ctx, notificationStore, notifier) - - var catalog metastore.Catalog - // TODO: move this to server.go - if db == nil { - catalog = coordinator.NewMemoryCatalogWithNotification(notificationStore) - } else { - txnImpl := dbcore.NewTxImpl() - metaDomain := dao.NewMetaDomain() - catalog = coordinator.NewTableCatalogWithNotification(txnImpl, metaDomain, notificationStore) - } - meta, err := NewMetaTable(s.ctx, catalog) - if err != nil { - return nil, err - } - meta.SetNotificationProcessor(notificationProcessor) - - s.meta = meta - s.notificationProcessor = notificationProcessor - - return s, nil -} - -func (s *Coordinator) Start() error { - err := s.notificationProcessor.Start() - if err != nil { - log.Printf("Failed to start notification processor: %v", err) - return err - } - return nil -} - -func (s *Coordinator) Stop() error { - err := s.notificationProcessor.Stop() - if err != nil { - log.Printf("Failed to stop notification processor: %v", err) - } - return nil -} - -func (c *Coordinator) assignCollection(collectionID types.UniqueID) (string, error) { - return c.collectionAssignmentPolicy.AssignCollection(collectionID) -} diff --git a/go/coordinator/internal/coordinator/meta.go b/go/coordinator/internal/coordinator/meta.go deleted file mode 100644 index f6f2df7584e..00000000000 --- a/go/coordinator/internal/coordinator/meta.go +++ /dev/null @@ -1,414 +0,0 @@ -package coordinator - -import ( - "context" - "sync" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/notification" - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/pingcap/log" - "go.uber.org/zap" -) - -// IMeta is an interface that defines methods for the cache of the catalog. -type IMeta interface { - ResetState(ctx context.Context) error - AddCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error) - GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*model.Collection, error) - DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error - UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection) (*model.Collection, error) - AddSegment(ctx context.Context, createSegment *model.CreateSegment) error - GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*model.Segment, error) - DeleteSegment(ctx context.Context, segmentID types.UniqueID) error - UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment) (*model.Segment, error) - CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error) - GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error) - CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) - GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error) - SetNotificationProcessor(notificationProcessor notification.NotificationProcessor) -} - -// MetaTable is an implementation of IMeta. It loads the system catalog during startup -// and caches in memory. The implmentation needs to make sure that the in memory cache -// is consistent with the system catalog. -// -// Operations of MetaTable are protected by a read write lock and are thread safe. -type MetaTable struct { - ddLock sync.RWMutex - ctx context.Context - catalog metastore.Catalog - segmentsCache map[types.UniqueID]*model.Segment - tenantDatabaseCollectionCache map[string]map[string]map[types.UniqueID]*model.Collection - tenantDatabaseCache map[string]map[string]*model.Database - notificationProcessor notification.NotificationProcessor -} - -var _ IMeta = (*MetaTable)(nil) - -func NewMetaTable(ctx context.Context, catalog metastore.Catalog) (*MetaTable, error) { - mt := &MetaTable{ - ctx: ctx, - catalog: catalog, - segmentsCache: make(map[types.UniqueID]*model.Segment), - tenantDatabaseCollectionCache: make(map[string]map[string]map[types.UniqueID]*model.Collection), - tenantDatabaseCache: make(map[string]map[string]*model.Database), - } - if err := mt.reloadWithLock(); err != nil { - return nil, err - } - return mt, nil -} - -func (mt *MetaTable) reloadWithLock() error { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - return mt.reload() -} - -func (mt *MetaTable) reload() error { - tenants, err := mt.catalog.GetAllTenants(mt.ctx, 0) - if err != nil { - return err - } - for _, tenant := range tenants { - tenantID := tenant.Name - mt.tenantDatabaseCollectionCache[tenantID] = make(map[string]map[types.UniqueID]*model.Collection) - mt.tenantDatabaseCache[tenantID] = make(map[string]*model.Database) - } - // reload databases - databases, err := mt.catalog.GetAllDatabases(mt.ctx, 0) - if err != nil { - return err - } - for _, database := range databases { - databaseName := database.Name - tenantID := database.Tenant - mt.tenantDatabaseCollectionCache[tenantID][databaseName] = make(map[types.UniqueID]*model.Collection) - mt.tenantDatabaseCache[tenantID][databaseName] = database - } - for tenantID, databases := range mt.tenantDatabaseCollectionCache { - for databaseName := range databases { - collections, err := mt.catalog.GetCollections(mt.ctx, types.NilUniqueID(), nil, nil, tenantID, databaseName) - if err != nil { - return err - } - for _, collection := range collections { - mt.tenantDatabaseCollectionCache[tenantID][databaseName][collection.ID] = collection - } - } - } - - oldSegments, err := mt.catalog.GetSegments(mt.ctx, types.NilUniqueID(), nil, nil, nil, types.NilUniqueID(), 0) - if err != nil { - return err - } - // reload is idempotent - mt.segmentsCache = make(map[types.UniqueID]*model.Segment) - for _, segment := range oldSegments { - mt.segmentsCache[segment.ID] = segment - } - return nil -} - -func (mt *MetaTable) SetNotificationProcessor(notificationProcessor notification.NotificationProcessor) { - mt.notificationProcessor = notificationProcessor -} - -func (mt *MetaTable) ResetState(ctx context.Context) error { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - if err := mt.catalog.ResetState(ctx); err != nil { - return err - } - mt.segmentsCache = make(map[types.UniqueID]*model.Segment) - mt.tenantDatabaseCache = make(map[string]map[string]*model.Database) - mt.tenantDatabaseCollectionCache = make(map[string]map[string]map[types.UniqueID]*model.Collection) - - if err := mt.reload(); err != nil { - return err - } - return nil -} - -func (mt *MetaTable) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error) { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - tenant := createDatabase.Tenant - databaseName := createDatabase.Name - if _, ok := mt.tenantDatabaseCache[tenant]; !ok { - log.Error("tenant not found", zap.Any("tenant", tenant)) - return nil, common.ErrTenantNotFound - } - if _, ok := mt.tenantDatabaseCache[tenant][databaseName]; ok { - log.Error("database already exists", zap.Any("database", databaseName)) - return nil, common.ErrDatabaseUniqueConstraintViolation - } - database, err := mt.catalog.CreateDatabase(ctx, createDatabase, createDatabase.Ts) - if err != nil { - log.Info("create database failed", zap.Error(err)) - return nil, err - } - mt.tenantDatabaseCache[tenant][databaseName] = database - mt.tenantDatabaseCollectionCache[tenant][databaseName] = make(map[types.UniqueID]*model.Collection) - return database, nil -} - -func (mt *MetaTable) GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error) { - mt.ddLock.RLock() - defer mt.ddLock.RUnlock() - - tenant := getDatabase.Tenant - databaseName := getDatabase.Name - if _, ok := mt.tenantDatabaseCache[tenant]; !ok { - log.Error("tenant not found", zap.Any("tenant", tenant)) - return nil, common.ErrTenantNotFound - } - if _, ok := mt.tenantDatabaseCache[tenant][databaseName]; !ok { - log.Error("database not found", zap.Any("database", databaseName)) - return nil, common.ErrDatabaseNotFound - } - - return mt.tenantDatabaseCache[tenant][databaseName], nil -} - -func (mt *MetaTable) CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - tenantName := createTenant.Name - if _, ok := mt.tenantDatabaseCache[tenantName]; ok { - log.Error("tenant already exists", zap.Any("tenant", tenantName)) - return nil, common.ErrTenantUniqueConstraintViolation - } - tenant, err := mt.catalog.CreateTenant(ctx, createTenant, createTenant.Ts) - if err != nil { - return nil, err - } - mt.tenantDatabaseCache[tenantName] = make(map[string]*model.Database) - mt.tenantDatabaseCollectionCache[tenantName] = make(map[string]map[types.UniqueID]*model.Collection) - return tenant, nil -} - -func (mt *MetaTable) GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error) { - mt.ddLock.RLock() - defer mt.ddLock.RUnlock() - tenantID := getTenant.Name - if _, ok := mt.tenantDatabaseCache[tenantID]; !ok { - log.Error("tenant not found", zap.Any("tenant", tenantID)) - return nil, common.ErrTenantNotFound - } - return &model.Tenant{Name: tenantID}, nil -} - -func (mt *MetaTable) AddCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error) { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - tenantID := createCollection.TenantID - databaseName := createCollection.DatabaseName - if _, ok := mt.tenantDatabaseCollectionCache[tenantID]; !ok { - log.Error("tenant not found", zap.Any("tenantID", tenantID)) - return nil, common.ErrTenantNotFound - } - if _, ok := mt.tenantDatabaseCollectionCache[tenantID][databaseName]; !ok { - log.Error("database not found", zap.Any("databaseName", databaseName)) - return nil, common.ErrDatabaseNotFound - } - collection, err := mt.catalog.CreateCollection(ctx, createCollection, createCollection.Ts) - if err != nil { - log.Error("create collection failed", zap.Error(err)) - return nil, err - } - mt.tenantDatabaseCollectionCache[tenantID][databaseName][collection.ID] = collection - log.Info("collection added", zap.Any("collection", mt.tenantDatabaseCollectionCache[tenantID][databaseName][collection.ID])) - - triggerMessage := notification.TriggerMessage{ - Msg: model.Notification{ - CollectionID: collection.ID.String(), - Type: model.NotificationTypeCreateCollection, - Status: model.NotificationStatusPending, - }, - ResultChan: make(chan error), - } - mt.notificationProcessor.Trigger(ctx, triggerMessage) - return collection, nil -} - -func (mt *MetaTable) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*model.Collection, error) { - mt.ddLock.RLock() - defer mt.ddLock.RUnlock() - - // There are three cases - // In the case of getting by id, we do not care about the tenant and database name. - // In the case of getting by name, we need the fully qualified path of the collection which is the tenant/database/name. - // In the case of getting by topic, we need the fully qualified path of the collection which is the tenant/database/topic. - collections := make([]*model.Collection, 0, len(mt.tenantDatabaseCollectionCache)) - if collectionID != types.NilUniqueID() { - // Case 1: getting by id - // Due to how the cache is constructed, we iterate over the whole cache to find the collection. - // This is not efficient but it is not a problem for now because the number of collections is small. - // HACK warning. TODO: fix this when we remove the cache. - for _, search_databases := range mt.tenantDatabaseCollectionCache { - for _, search_collections := range search_databases { - for _, collection := range search_collections { - if model.FilterCollection(collection, collectionID, collectionName, collectionTopic) { - collections = append(collections, collection) - } - } - } - } - } else { - // Case 2 & 3: getting by name or topic - // Note: The support for case 3 is not correct here, we shouldn't require the database name and tenant to get by topic. - if _, ok := mt.tenantDatabaseCollectionCache[tenantID]; !ok { - log.Error("tenant not found", zap.Any("tenantID", tenantID)) - return nil, common.ErrTenantNotFound - } - if _, ok := mt.tenantDatabaseCollectionCache[tenantID][databaseName]; !ok { - return nil, common.ErrDatabaseNotFound - } - for _, collection := range mt.tenantDatabaseCollectionCache[tenantID][databaseName] { - if model.FilterCollection(collection, collectionID, collectionName, collectionTopic) { - collections = append(collections, collection) - } - } - } - log.Info("meta collections", zap.Any("collections", collections)) - return collections, nil - -} - -func (mt *MetaTable) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - tenantID := deleteCollection.TenantID - databaseName := deleteCollection.DatabaseName - collectionID := deleteCollection.ID - if _, ok := mt.tenantDatabaseCollectionCache[tenantID]; !ok { - log.Error("tenant not found", zap.Any("tenantID", tenantID)) - return common.ErrTenantNotFound - } - if _, ok := mt.tenantDatabaseCollectionCache[tenantID][databaseName]; !ok { - log.Error("database not found", zap.Any("databaseName", databaseName)) - return common.ErrDatabaseNotFound - } - collections := mt.tenantDatabaseCollectionCache[tenantID][databaseName] - - if _, ok := collections[collectionID]; !ok { - log.Error("collection not found", zap.Any("collectionID", collectionID)) - return common.ErrCollectionDeleteNonExistingCollection - } - - if err := mt.catalog.DeleteCollection(ctx, deleteCollection); err != nil { - return err - } - delete(collections, collectionID) - log.Info("collection deleted", zap.Any("collection", deleteCollection)) - - triggerMessage := notification.TriggerMessage{ - Msg: model.Notification{ - CollectionID: collectionID.String(), - Type: model.NotificationTypeDeleteCollection, - Status: model.NotificationStatusPending, - }, - ResultChan: make(chan error), - } - mt.notificationProcessor.Trigger(ctx, triggerMessage) - return nil -} - -func (mt *MetaTable) UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection) (*model.Collection, error) { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - var oldCollection *model.Collection - for tenant := range mt.tenantDatabaseCollectionCache { - for database := range mt.tenantDatabaseCollectionCache[tenant] { - for _, collection := range mt.tenantDatabaseCollectionCache[tenant][database] { - if collection.ID == updateCollection.ID { - oldCollection = collection - break - } - } - } - } - if oldCollection == nil { - log.Error("collection not found", zap.Any("collectionID", updateCollection.ID)) - return nil, common.ErrCollectionNotFound - } - - updateCollection.DatabaseName = oldCollection.DatabaseName - updateCollection.TenantID = oldCollection.TenantID - - collection, err := mt.catalog.UpdateCollection(ctx, updateCollection, updateCollection.Ts) - if err != nil { - return nil, err - } - mt.tenantDatabaseCollectionCache[collection.TenantID][collection.DatabaseName][collection.ID] = collection - log.Info("collection updated", zap.Any("collection", collection)) - return collection, nil -} - -func (mt *MetaTable) AddSegment(ctx context.Context, createSegment *model.CreateSegment) error { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - segment, err := mt.catalog.CreateSegment(ctx, createSegment, createSegment.Ts) - if err != nil { - return err - } - mt.segmentsCache[createSegment.ID] = segment - log.Info("segment added", zap.Any("segment", segment)) - return nil -} - -func (mt *MetaTable) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*model.Segment, error) { - mt.ddLock.RLock() - defer mt.ddLock.RUnlock() - - segments := make([]*model.Segment, 0, len(mt.segmentsCache)) - for _, segment := range mt.segmentsCache { - if model.FilterSegments(segment, segmentID, segmentType, scope, topic, collectionID) { - segments = append(segments, segment) - } - } - log.Info("meta get segments", zap.Any("segments", segments)) - return segments, nil -} - -func (mt *MetaTable) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - if _, ok := mt.segmentsCache[segmentID]; !ok { - return common.ErrSegmentDeleteNonExistingSegment - } - - if err := mt.catalog.DeleteSegment(ctx, segmentID); err != nil { - log.Error("delete segment failed", zap.Error(err)) - return err - } - delete(mt.segmentsCache, segmentID) - log.Info("segment deleted", zap.Any("segmentID", segmentID)) - return nil -} - -func (mt *MetaTable) UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment) (*model.Segment, error) { - mt.ddLock.Lock() - defer mt.ddLock.Unlock() - - segment, err := mt.catalog.UpdateSegment(ctx, updateSegment, updateSegment.Ts) - if err != nil { - log.Error("update segment failed", zap.Error(err)) - return nil, err - } - mt.segmentsCache[updateSegment.ID] = segment - log.Info("segment updated", zap.Any("segment", segment)) - return segment, nil -} diff --git a/go/coordinator/internal/coordinator/meta_test.go b/go/coordinator/internal/coordinator/meta_test.go deleted file mode 100644 index d40ddf4ea33..00000000000 --- a/go/coordinator/internal/coordinator/meta_test.go +++ /dev/null @@ -1,94 +0,0 @@ -package coordinator - -import ( - "context" - "testing" - - "github.com/chroma/chroma-coordinator/internal/metastore/coordinator" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" - "pgregory.net/rapid" -) - -func testMeta(t *rapid.T) { - catalog := coordinator.NewMemoryCatalog() - mt, err := NewMetaTable(context.Background(), catalog) - if err != nil { - t.Fatalf("error creating meta table: %v", err) - } - t.Repeat(map[string]func(*rapid.T){ - "generate_collection": func(t *rapid.T) { - collection := rapid.Custom[*model.CreateCollection](func(t *rapid.T) *model.CreateCollection { - return &model.CreateCollection{ - ID: genCollectinID(t), - Name: rapid.String().Draw(t, "name"), - // Dimension: rapid.Int32().Draw(t, "dimension"), - Metadata: rapid.Custom[*model.CollectionMetadata[model.CollectionMetadataValueType]](func(t *rapid.T) *model.CollectionMetadata[model.CollectionMetadataValueType] { - return &model.CollectionMetadata[model.CollectionMetadataValueType]{ - Metadata: rapid.MapOf[string, model.CollectionMetadataValueType](rapid.StringMatching(`[a-zA-Z0-9_]+`), drawMetadata(t)).Draw(t, "metadata"), - } - }).Draw(t, "metadata"), - } - }).Draw(t, "collection") - if _, err := mt.catalog.CreateCollection(context.Background(), collection, 0); err != nil { - t.Fatalf("error creating collection: %v", err) - } - }, - "reload": func(t *rapid.T) { - if err := mt.reload(); err != nil { - t.Fatalf("error reloading meta table: %v", err) - } - }, - "add_collection": func(t *rapid.T) { - if err := mt.reload(); err != nil { - t.Fatalf("error reloading meta table: %v", err) - } - collection := rapid.Custom[*model.CreateCollection](func(t *rapid.T) *model.CreateCollection { - return &model.CreateCollection{ - ID: genCollectinID(t), - Name: rapid.String().Draw(t, "name"), - //Dimension: rapid.Int32().Draw(t, "dimension"), - Metadata: rapid.Custom[*model.CollectionMetadata[model.CollectionMetadataValueType]](func(t *rapid.T) *model.CollectionMetadata[model.CollectionMetadataValueType] { - return &model.CollectionMetadata[model.CollectionMetadataValueType]{ - Metadata: rapid.MapOf[string, model.CollectionMetadataValueType](rapid.StringMatching(`[a-zA-Z0-9_]+`), drawMetadata(t)).Draw(t, "metadata"), - } - }).Draw(t, "metadata"), - } - }).Draw(t, "collection") - - if _, err := mt.AddCollection(context.Background(), collection); err != nil { - t.Fatalf("error adding collection: %v", err) - } - }, - }) -} - -func drawMetadata(t *rapid.T) *rapid.Generator[model.CollectionMetadataValueType] { - return rapid.OneOf[model.CollectionMetadataValueType]( - rapid.Custom[model.CollectionMetadataValueType](func(t *rapid.T) model.CollectionMetadataValueType { - return &model.CollectionMetadataValueStringType{ - Value: rapid.String().Draw(t, "string_value"), - } - }), - rapid.Custom[model.CollectionMetadataValueType](func(t *rapid.T) model.CollectionMetadataValueType { - return &model.CollectionMetadataValueInt64Type{ - Value: rapid.Int64().Draw(t, "int_value"), - } - }), - rapid.Custom[model.CollectionMetadataValueType](func(t *rapid.T) model.CollectionMetadataValueType { - return &model.CollectionMetadataValueFloat64Type{ - Value: rapid.Float64().Draw(t, "float_value"), - } - }), - ) -} - -func genCollectinID(t *rapid.T) types.UniqueID { - return rapid.Custom[types.UniqueID](func(t *rapid.T) types.UniqueID { - return types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "uuid")) - }).Draw(t, "collection_id") -} - -func TestMeta(t *testing.T) { - // rapid.Check(t, testMeta) -} diff --git a/go/coordinator/internal/grpccoordinator/collection_service_test.go b/go/coordinator/internal/grpccoordinator/collection_service_test.go deleted file mode 100644 index 390b08f7607..00000000000 --- a/go/coordinator/internal/grpccoordinator/collection_service_test.go +++ /dev/null @@ -1,125 +0,0 @@ -package grpccoordinator - -import ( - "context" - "testing" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/grpccoordinator/grpcutils" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbcore" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" - "pgregory.net/rapid" -) - -// CreateCollection -// Collection created successfully are visible to ListCollections -// Collection created should have the right metadata, the metadata should be a flat map, with keys as strings and values as strings, ints, or floats -// Collection created should have the right name -// Collection created should have the right ID -// Collection created should have the right topic -// Collection created should have the right timestamp -func testCollection(t *rapid.T) { - db := dbcore.ConfigDatabaseForTesting() - s, err := NewWithGrpcProvider(Config{ - AssignmentPolicy: "simple", - SystemCatalogProvider: "memory", - NotificationStoreProvider: "memory", - NotifierProvider: "memory", - Testing: true}, grpcutils.Default, db) - if err != nil { - t.Fatalf("error creating server: %v", err) - } - var state []*coordinatorpb.Collection - var collectionsWithErrors []*coordinatorpb.Collection - - t.Repeat(map[string]func(*rapid.T){ - "create_collection": func(t *rapid.T) { - stringValue := generateStringMetadataValue(t) - intValue := generateInt64MetadataValue(t) - floatValue := generateFloat64MetadataValue(t) - getOrCreate := false - - createCollectionRequest := rapid.Custom[*coordinatorpb.CreateCollectionRequest](func(t *rapid.T) *coordinatorpb.CreateCollectionRequest { - return &coordinatorpb.CreateCollectionRequest{ - Id: rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "collection_id"), - Name: rapid.String().Draw(t, "collection_name"), - Metadata: &coordinatorpb.UpdateMetadata{ - Metadata: map[string]*coordinatorpb.UpdateMetadataValue{ - "string_value": stringValue, - "int_value": intValue, - "float_value": floatValue, - }, - }, - GetOrCreate: &getOrCreate, - } - }).Draw(t, "create_collection_request") - - ctx := context.Background() - res, err := s.CreateCollection(ctx, createCollectionRequest) - if err != nil { - if err == common.ErrCollectionNameEmpty && createCollectionRequest.Name == "" { - t.Logf("expected error for empty collection name") - collectionsWithErrors = append(collectionsWithErrors, res.Collection) - } else if err == common.ErrCollectionTopicEmpty { - t.Logf("expected error for empty collection topic") - collectionsWithErrors = append(collectionsWithErrors, res.Collection) - // TODO: check the topic name not empty - } else { - t.Fatalf("error creating collection: %v", err) - collectionsWithErrors = append(collectionsWithErrors, res.Collection) - } - } - - getCollectionsRequest := coordinatorpb.GetCollectionsRequest{ - Id: &createCollectionRequest.Id, - } - if err == nil { - // verify the correctness - GetCollectionsResponse, err := s.GetCollections(ctx, &getCollectionsRequest) - if err != nil { - t.Fatalf("error getting collections: %v", err) - } - collectionList := GetCollectionsResponse.GetCollections() - if len(collectionList) != 1 { - t.Fatalf("More than 1 collection with the same collection id") - } - for _, collection := range collectionList { - if collection.Id != createCollectionRequest.Id { - t.Fatalf("collection id is the right value") - } - } - state = append(state, res.Collection) - } - }, - "get_collections": func(t *rapid.T) { - }, - }) -} - -func generateStringMetadataValue(t *rapid.T) *coordinatorpb.UpdateMetadataValue { - return &coordinatorpb.UpdateMetadataValue{ - Value: &coordinatorpb.UpdateMetadataValue_StringValue{ - StringValue: rapid.String().Draw(t, "string_value"), - }, - } -} - -func generateInt64MetadataValue(t *rapid.T) *coordinatorpb.UpdateMetadataValue { - return &coordinatorpb.UpdateMetadataValue{ - Value: &coordinatorpb.UpdateMetadataValue_IntValue{ - IntValue: rapid.Int64().Draw(t, "int_value"), - }, - } -} - -func generateFloat64MetadataValue(t *rapid.T) *coordinatorpb.UpdateMetadataValue { - return &coordinatorpb.UpdateMetadataValue{ - Value: &coordinatorpb.UpdateMetadataValue_FloatValue{ - FloatValue: rapid.Float64().Draw(t, "float_value"), - }, - } -} - -func TestCollection(t *testing.T) { - // rapid.Check(t, testCollection) -} diff --git a/go/coordinator/internal/metastore/coordinator/memory_catalog.go b/go/coordinator/internal/metastore/coordinator/memory_catalog.go deleted file mode 100644 index 439911cb754..00000000000 --- a/go/coordinator/internal/metastore/coordinator/memory_catalog.go +++ /dev/null @@ -1,370 +0,0 @@ -package coordinator - -import ( - "context" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/notification" - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/pingcap/log" - "go.uber.org/zap" -) - -// MemoryCatalog is a reference implementation of Catalog interface to ensure the -// application logic is correctly implemented. -type MemoryCatalog struct { - segments map[types.UniqueID]*model.Segment - tenantDatabaseCollections map[string]map[string]map[types.UniqueID]*model.Collection - tenantDatabases map[string]map[string]*model.Database - store notification.NotificationStore -} - -var _ metastore.Catalog = (*MemoryCatalog)(nil) - -func NewMemoryCatalog() *MemoryCatalog { - memoryCatalog := MemoryCatalog{ - segments: make(map[types.UniqueID]*model.Segment), - tenantDatabaseCollections: make(map[string]map[string]map[types.UniqueID]*model.Collection), - tenantDatabases: make(map[string]map[string]*model.Database), - } - // Add a default tenant and database - memoryCatalog.tenantDatabases[common.DefaultTenant] = make(map[string]*model.Database) - memoryCatalog.tenantDatabases[common.DefaultTenant][common.DefaultDatabase] = &model.Database{ - ID: types.NilUniqueID().String(), - Name: common.DefaultDatabase, - Tenant: common.DefaultTenant, - } - memoryCatalog.tenantDatabaseCollections[common.DefaultTenant] = make(map[string]map[types.UniqueID]*model.Collection) - memoryCatalog.tenantDatabaseCollections[common.DefaultTenant][common.DefaultDatabase] = make(map[types.UniqueID]*model.Collection) - return &memoryCatalog -} - -func NewMemoryCatalogWithNotification(store notification.NotificationStore) *MemoryCatalog { - memoryCatalog := NewMemoryCatalog() - memoryCatalog.store = store - return memoryCatalog -} - -func (mc *MemoryCatalog) ResetState(ctx context.Context) error { - mc.segments = make(map[types.UniqueID]*model.Segment) - mc.tenantDatabases = make(map[string]map[string]*model.Database) - mc.tenantDatabases[common.DefaultTenant] = make(map[string]*model.Database) - mc.tenantDatabases[common.DefaultTenant][common.DefaultDatabase] = &model.Database{ - ID: types.NilUniqueID().String(), - Name: common.DefaultDatabase, - Tenant: common.DefaultTenant, - } - mc.tenantDatabaseCollections[common.DefaultTenant] = make(map[string]map[types.UniqueID]*model.Collection) - mc.tenantDatabaseCollections[common.DefaultTenant][common.DefaultDatabase] = make(map[types.UniqueID]*model.Collection) - return nil -} - -func (mc *MemoryCatalog) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase, ts types.Timestamp) (*model.Database, error) { - tenant := createDatabase.Tenant - databaseName := createDatabase.Name - if _, ok := mc.tenantDatabases[tenant]; !ok { - log.Error("tenant not found", zap.String("tenant", tenant)) - return nil, common.ErrTenantNotFound - } - if _, ok := mc.tenantDatabases[tenant][databaseName]; ok { - log.Error("database already exists", zap.String("database", databaseName)) - return nil, common.ErrDatabaseUniqueConstraintViolation - } - mc.tenantDatabases[tenant][databaseName] = &model.Database{ - ID: createDatabase.ID, - Name: createDatabase.Name, - Tenant: createDatabase.Tenant, - } - mc.tenantDatabaseCollections[tenant][databaseName] = make(map[types.UniqueID]*model.Collection) - log.Info("database created", zap.Any("database", mc.tenantDatabases[tenant][databaseName])) - return mc.tenantDatabases[tenant][databaseName], nil -} - -func (mc *MemoryCatalog) GetDatabases(ctx context.Context, getDatabase *model.GetDatabase, ts types.Timestamp) (*model.Database, error) { - tenant := getDatabase.Tenant - databaseName := getDatabase.Name - if _, ok := mc.tenantDatabases[tenant]; !ok { - log.Error("tenant not found", zap.String("tenant", tenant)) - return nil, common.ErrTenantNotFound - } - if _, ok := mc.tenantDatabases[tenant][databaseName]; !ok { - log.Error("database not found", zap.String("database", databaseName)) - return nil, common.ErrDatabaseNotFound - } - log.Info("database found", zap.Any("database", mc.tenantDatabases[tenant][databaseName])) - return mc.tenantDatabases[tenant][databaseName], nil -} - -func (mc *MemoryCatalog) GetAllDatabases(ctx context.Context, ts types.Timestamp) ([]*model.Database, error) { - databases := make([]*model.Database, 0) - for _, database := range mc.tenantDatabases { - for _, db := range database { - databases = append(databases, db) - } - } - return databases, nil -} - -func (mc *MemoryCatalog) CreateTenant(ctx context.Context, createTenant *model.CreateTenant, ts types.Timestamp) (*model.Tenant, error) { - tenant := createTenant.Name - if _, ok := mc.tenantDatabases[tenant]; ok { - log.Error("tenant already exists", zap.String("tenant", tenant)) - return nil, common.ErrTenantUniqueConstraintViolation - } - mc.tenantDatabases[tenant] = make(map[string]*model.Database) - mc.tenantDatabaseCollections[tenant] = make(map[string]map[types.UniqueID]*model.Collection) - return &model.Tenant{Name: tenant}, nil -} - -func (mc *MemoryCatalog) GetTenants(ctx context.Context, getTenant *model.GetTenant, ts types.Timestamp) (*model.Tenant, error) { - tenant := getTenant.Name - if _, ok := mc.tenantDatabases[tenant]; !ok { - log.Error("tenant not found", zap.String("tenant", tenant)) - return nil, common.ErrTenantNotFound - } - return &model.Tenant{Name: tenant}, nil -} - -func (mc *MemoryCatalog) GetAllTenants(ctx context.Context, ts types.Timestamp) ([]*model.Tenant, error) { - tenants := make([]*model.Tenant, 0, len(mc.tenantDatabases)) - for tenant := range mc.tenantDatabases { - tenants = append(tenants, &model.Tenant{Name: tenant}) - } - return tenants, nil -} - -func (mc *MemoryCatalog) CreateCollection(ctx context.Context, createCollection *model.CreateCollection, ts types.Timestamp) (*model.Collection, error) { - collectionName := createCollection.Name - tenantID := createCollection.TenantID - databaseName := createCollection.DatabaseName - - if _, ok := mc.tenantDatabaseCollections[tenantID]; !ok { - log.Error("tenant not found", zap.String("tenant", tenantID)) - return nil, common.ErrTenantNotFound - } - if _, ok := mc.tenantDatabaseCollections[tenantID][databaseName]; !ok { - log.Error("database not found", zap.String("database", databaseName)) - return nil, common.ErrDatabaseNotFound - } - // Check if the collection already by global id - for tenant := range mc.tenantDatabaseCollections { - for database := range mc.tenantDatabaseCollections[tenant] { - collections := mc.tenantDatabaseCollections[tenant][database] - if _, ok := collections[createCollection.ID]; ok { - if tenant != tenantID || database != databaseName { - log.Info("collection already exists", zap.Any("collection", collections[createCollection.ID])) - return nil, common.ErrCollectionUniqueConstraintViolation - } else if !createCollection.GetOrCreate { - return nil, common.ErrCollectionUniqueConstraintViolation - } - } - - } - } - // Check if the collection already exists in database by colllection name - collections := mc.tenantDatabaseCollections[tenantID][databaseName] - for _, collection := range collections { - if collection.Name == collectionName { - log.Info("collection already exists", zap.Any("collection", collections[createCollection.ID])) - if createCollection.GetOrCreate { - if createCollection.Metadata != nil { - // For getOrCreate, update the metadata - collection.Metadata = createCollection.Metadata - } - return collection, nil - } else { - return nil, common.ErrCollectionUniqueConstraintViolation - } - } - } - collection := &model.Collection{ - ID: createCollection.ID, - Name: createCollection.Name, - Topic: createCollection.Topic, - Dimension: createCollection.Dimension, - Metadata: createCollection.Metadata, - Created: true, - TenantID: createCollection.TenantID, - DatabaseName: createCollection.DatabaseName, - } - log.Info("collection created", zap.Any("collection", collection)) - collections[collection.ID] = collection - return collection, nil -} - -func (mc *MemoryCatalog) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*model.Collection, error) { - if _, ok := mc.tenantDatabaseCollections[tenantID]; !ok { - log.Error("tenant not found", zap.String("tenant", tenantID)) - return nil, common.ErrTenantNotFound - } - if _, ok := mc.tenantDatabaseCollections[tenantID][databaseName]; !ok { - log.Error("database not found", zap.String("database", databaseName)) - return nil, common.ErrDatabaseNotFound - } - collections := make([]*model.Collection, 0, len(mc.tenantDatabaseCollections[tenantID][databaseName])) - for _, collection := range mc.tenantDatabaseCollections[tenantID][databaseName] { - if model.FilterCollection(collection, collectionID, collectionName, collectionTopic) { - collections = append(collections, collection) - } - } - return collections, nil -} - -func (mc *MemoryCatalog) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error { - tenantID := deleteCollection.TenantID - databaseName := deleteCollection.DatabaseName - collectionID := deleteCollection.ID - if _, ok := mc.tenantDatabaseCollections[tenantID]; !ok { - log.Error("tenant not found", zap.String("tenant", tenantID)) - return common.ErrTenantNotFound - } - if _, ok := mc.tenantDatabaseCollections[tenantID][databaseName]; !ok { - log.Error("database not found", zap.String("database", databaseName)) - return common.ErrDatabaseNotFound - } - collections := mc.tenantDatabaseCollections[tenantID][databaseName] - if _, ok := collections[collectionID]; !ok { - log.Error("collection not found", zap.String("collection", collectionID.String())) - return common.ErrCollectionDeleteNonExistingCollection - } - delete(collections, collectionID) - log.Info("collection deleted", zap.String("collection", collectionID.String())) - mc.store.AddNotification(ctx, model.Notification{ - CollectionID: collectionID.String(), - Type: model.NotificationTypeDeleteCollection, - Status: model.NotificationStatusPending, - }) - return nil -} - -func (mc *MemoryCatalog) UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection, ts types.Timestamp) (*model.Collection, error) { - collectionID := updateCollection.ID - var oldCollection *model.Collection - for tenant := range mc.tenantDatabaseCollections { - for database := range mc.tenantDatabaseCollections[tenant] { - log.Info("database", zap.Any("database", database)) - collections := mc.tenantDatabaseCollections[tenant][database] - if _, ok := collections[collectionID]; ok { - oldCollection = collections[collectionID] - } - } - } - - topic := updateCollection.Topic - if topic != nil { - oldCollection.Topic = *topic - } - name := updateCollection.Name - if name != nil { - oldCollection.Name = *name - } - if updateCollection.Dimension != nil { - oldCollection.Dimension = updateCollection.Dimension - } - - // Case 1: if resetMetadata is true, then delete all metadata for the collection - // Case 2: if resetMetadata is true and metadata is not nil -> THIS SHOULD NEVER HAPPEN - // Case 3: if resetMetadata is false, and the metadata is not nil - set the metadata to the value in metadata - // Case 4: if resetMetadata is false and metadata is nil, then leave the metadata as is - resetMetadata := updateCollection.ResetMetadata - if resetMetadata { - oldCollection.Metadata = nil - } else { - if updateCollection.Metadata != nil { - oldCollection.Metadata = updateCollection.Metadata - } - } - tenantID := oldCollection.TenantID - databaseName := oldCollection.DatabaseName - mc.tenantDatabaseCollections[tenantID][databaseName][oldCollection.ID] = oldCollection - // Better to return a copy of the collection to avoid being modified by others. - log.Debug("collection metadata", zap.Any("metadata", oldCollection.Metadata)) - return oldCollection, nil -} - -func (mc *MemoryCatalog) CreateSegment(ctx context.Context, createSegment *model.CreateSegment, ts types.Timestamp) (*model.Segment, error) { - if _, ok := mc.segments[createSegment.ID]; ok { - return nil, common.ErrSegmentUniqueConstraintViolation - } - - segment := &model.Segment{ - ID: createSegment.ID, - Topic: createSegment.Topic, - Type: createSegment.Type, - Scope: createSegment.Scope, - CollectionID: createSegment.CollectionID, - Metadata: createSegment.Metadata, - } - mc.segments[createSegment.ID] = segment - log.Debug("segment created", zap.Any("segment", segment)) - return segment, nil -} - -func (mc *MemoryCatalog) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID, ts types.Timestamp) ([]*model.Segment, error) { - segments := make([]*model.Segment, 0, len(mc.segments)) - for _, segment := range mc.segments { - if model.FilterSegments(segment, segmentID, segmentType, scope, topic, collectionID) { - segments = append(segments, segment) - } - } - return segments, nil -} - -func (mc *MemoryCatalog) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error { - if _, ok := mc.segments[segmentID]; !ok { - return common.ErrSegmentDeleteNonExistingSegment - } - - delete(mc.segments, segmentID) - return nil -} - -func (mc *MemoryCatalog) UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment, ts types.Timestamp) (*model.Segment, error) { - // Case 1: if ResetTopic is true and topic is nil, then set the topic to nil - // Case 2: if ResetTopic is true and topic is not nil -> THIS SHOULD NEVER HAPPEN - // Case 3: if ResetTopic is false and topic is not nil - set the topic to the value in topic - // Case 4: if ResetTopic is false and topic is nil, then leave the topic as is - oldSegment := mc.segments[updateSegment.ID] - topic := updateSegment.Topic - if updateSegment.ResetTopic { - if topic == nil { - oldSegment.Topic = nil - } - } else { - if topic != nil { - oldSegment.Topic = topic - } - } - collection := updateSegment.Collection - if updateSegment.ResetCollection { - if collection == nil { - oldSegment.CollectionID = types.NilUniqueID() - } - } else { - if collection != nil { - parsedCollectionID, err := types.ToUniqueID(collection) - if err != nil { - return nil, err - } - oldSegment.CollectionID = parsedCollectionID - } - } - resetMetadata := updateSegment.ResetMetadata - if resetMetadata { - oldSegment.Metadata = nil - } else { - if updateSegment.Metadata != nil { - for key, value := range updateSegment.Metadata.Metadata { - if value == nil { - oldSegment.Metadata.Remove(key) - } else { - oldSegment.Metadata.Set(key, value) - } - } - } - } - mc.segments[updateSegment.ID] = oldSegment - return oldSegment, nil -} diff --git a/go/coordinator/internal/metastore/coordinator/memory_catalog_test.go b/go/coordinator/internal/metastore/coordinator/memory_catalog_test.go deleted file mode 100644 index c7f4b2d6040..00000000000 --- a/go/coordinator/internal/metastore/coordinator/memory_catalog_test.go +++ /dev/null @@ -1,138 +0,0 @@ -package coordinator - -import ( - "context" - "testing" - - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/notification" - "github.com/chroma/chroma-coordinator/internal/types" -) - -const ( - defaultTenant = "default_tenant" - defaultDatabase = "default_database" -) - -func TestMemoryCatalog(t *testing.T) { - ctx := context.Background() - store := notification.NewMemoryNotificationStore() - mc := NewMemoryCatalogWithNotification(store) - - // Test CreateCollection - coll := &model.CreateCollection{ - ID: types.NewUniqueID(), - Name: "test-collection-name", - // Topic: "test-collection-topic", - Metadata: &model.CollectionMetadata[model.CollectionMetadataValueType]{ - Metadata: map[string]model.CollectionMetadataValueType{ - "test-metadata-key": &model.CollectionMetadataValueStringType{Value: "test-metadata-value"}, - }, - }, - TenantID: defaultTenant, - DatabaseName: defaultDatabase, - } - collection, err := mc.CreateCollection(ctx, coll, types.Timestamp(0)) - if err != nil { - t.Fatalf("unexpected error creating collection: %v", err) - } - // Test GetCollections - collections, err := mc.GetCollections(ctx, coll.ID, &coll.Name, nil, defaultTenant, defaultDatabase) - if err != nil { - t.Fatalf("unexpected error getting collections: %v", err) - } - if len(collections) != 1 { - t.Fatalf("expected 1 collection, got %d", len(collections)) - } - if collections[0] != collection { - t.Fatalf("expected collection %+v, got %+v", coll, collections[0]) - } - - // Test DeleteCollection - deleteCollection := &model.DeleteCollection{ - ID: coll.ID, - DatabaseName: defaultDatabase, - TenantID: defaultTenant, - } - if err := mc.DeleteCollection(ctx, deleteCollection); err != nil { - t.Fatalf("unexpected error deleting collection: %v", err) - } - - // Test CreateSegment - testTopic := "test-segment-topic" - createSegment := &model.CreateSegment{ - ID: types.NewUniqueID(), - Type: "test-segment-type", - Scope: "test-segment-scope", - Topic: &testTopic, - CollectionID: coll.ID, - Metadata: &model.SegmentMetadata[model.SegmentMetadataValueType]{ - Metadata: map[string]model.SegmentMetadataValueType{ - "test-metadata-key": &model.SegmentMetadataValueStringType{Value: "test-metadata-value"}, - }, - }, - } - segment, err := mc.CreateSegment(ctx, createSegment, types.Timestamp(0)) - if err != nil { - t.Fatalf("unexpected error creating segment: %v", err) - } - if len(mc.segments) != 1 { - t.Fatalf("expected 1 segment, got %d", len(mc.segments)) - } - - if mc.segments[createSegment.ID] != segment { - t.Fatalf("expected segment with ID %q, got %+v", createSegment.ID, mc.segments[createSegment.ID]) - } - - // Test GetSegments - segments, err := mc.GetSegments(ctx, createSegment.ID, &createSegment.Type, &createSegment.Scope, createSegment.Topic, coll.ID, types.Timestamp(0)) - if err != nil { - t.Fatalf("unexpected error getting segments: %v", err) - } - if len(segments) != 1 { - t.Fatalf("expected 1 segment, got %d", len(segments)) - } - if segments[0] != segment { - t.Fatalf("expected segment %+v, got %+v", createSegment, segments[0]) - } - - // Test CreateCollection - coll = &model.CreateCollection{ - ID: types.NewUniqueID(), - Name: "test-collection-name", - // Topic: "test-collection-topic", - Metadata: &model.CollectionMetadata[model.CollectionMetadataValueType]{ - Metadata: map[string]model.CollectionMetadataValueType{ - "test-metadata-key": &model.CollectionMetadataValueStringType{Value: "test-metadata-value"}, - }, - }, - TenantID: defaultTenant, - DatabaseName: defaultDatabase, - } - collection, err = mc.CreateCollection(ctx, coll, types.Timestamp(0)) - if err != nil { - t.Fatalf("unexpected error creating collection: %v", err) - } - - // Test GetCollections - collections, err = mc.GetCollections(ctx, coll.ID, &coll.Name, nil, defaultTenant, defaultDatabase) - if err != nil { - t.Fatalf("unexpected error getting collections: %v", err) - } - if len(collections) != 1 { - t.Fatalf("expected 1 collection, got %d", len(collections)) - } - if collections[0] != collection { - t.Fatalf("expected collection %+v, got %+v", coll, collections[0]) - } - - // Test DeleteCollection - deleteCollection = &model.DeleteCollection{ - ID: coll.ID, - DatabaseName: defaultDatabase, - TenantID: defaultTenant, - } - if err := mc.DeleteCollection(ctx, deleteCollection); err != nil { - t.Fatalf("unexpected error deleting collection: %v", err) - } -} diff --git a/go/coordinator/internal/metastore/db/dao/collection_test.go b/go/coordinator/internal/metastore/db/dao/collection_test.go deleted file mode 100644 index 1c2da046ec0..00000000000 --- a/go/coordinator/internal/metastore/db/dao/collection_test.go +++ /dev/null @@ -1,95 +0,0 @@ -package dao - -import ( - "testing" - - "github.com/pingcap/log" - "go.uber.org/zap" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/stretchr/testify/assert" - "gorm.io/driver/sqlite" - "gorm.io/gorm" -) - -func TestCollectionDb_GetCollections(t *testing.T) { - db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) - assert.NoError(t, err) - - err = db.AutoMigrate(&dbmodel.Tenant{}, &dbmodel.Database{}, &dbmodel.Collection{}, &dbmodel.CollectionMetadata{}) - db.Model(&dbmodel.Tenant{}).Create(&dbmodel.Tenant{ - ID: common.DefaultTenant, - }) - - databaseID := types.NilUniqueID().String() - db.Model(&dbmodel.Database{}).Create(&dbmodel.Database{ - ID: databaseID, - Name: common.DefaultDatabase, - TenantID: common.DefaultTenant, - }) - - assert.NoError(t, err) - name := "test_name" - topic := "test_topic" - collection := &dbmodel.Collection{ - ID: types.NewUniqueID().String(), - Name: &name, - Topic: &topic, - DatabaseID: databaseID, - } - err = db.Create(collection).Error - assert.NoError(t, err) - - testKey := "test" - testValue := "test" - metadata := &dbmodel.CollectionMetadata{ - CollectionID: collection.ID, - Key: &testKey, - StrValue: &testValue, - } - err = db.Create(metadata).Error - assert.NoError(t, err) - - collectionDb := &collectionDb{ - db: db, - } - - query := db.Table("collections").Select("collections.id") - rows, err := query.Rows() - assert.NoError(t, err) - for rows.Next() { - var collectionID string - err = rows.Scan(&collectionID) - assert.NoError(t, err) - log.Info("collectionID", zap.String("collectionID", collectionID)) - } - collections, err := collectionDb.GetCollections(nil, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Len(t, collections, 1) - assert.Equal(t, collection.ID, collections[0].Collection.ID) - assert.Equal(t, collection.Name, collections[0].Collection.Name) - assert.Equal(t, collection.Topic, collections[0].Collection.Topic) - assert.Len(t, collections[0].CollectionMetadata, 1) - assert.Equal(t, metadata.Key, collections[0].CollectionMetadata[0].Key) - assert.Equal(t, metadata.StrValue, collections[0].CollectionMetadata[0].StrValue) - - // Test when filtering by ID - collections, err = collectionDb.GetCollections(nil, nil, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Len(t, collections, 1) - assert.Equal(t, collection.ID, collections[0].Collection.ID) - - // Test when filtering by name - collections, err = collectionDb.GetCollections(nil, collection.Name, nil, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Len(t, collections, 1) - assert.Equal(t, collection.ID, collections[0].Collection.ID) - - // Test when filtering by topic - collections, err = collectionDb.GetCollections(nil, nil, collection.Topic, common.DefaultTenant, common.DefaultDatabase) - assert.NoError(t, err) - assert.Len(t, collections, 1) - assert.Equal(t, collection.ID, collections[0].Collection.ID) -} diff --git a/go/coordinator/internal/metastore/db/dao/segment_test.go b/go/coordinator/internal/metastore/db/dao/segment_test.go deleted file mode 100644 index 34522869faa..00000000000 --- a/go/coordinator/internal/metastore/db/dao/segment_test.go +++ /dev/null @@ -1,89 +0,0 @@ -package dao - -import ( - "testing" - - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/stretchr/testify/assert" - "gorm.io/driver/sqlite" - "gorm.io/gorm" -) - -func TestSegmentDb_GetSegments(t *testing.T) { - db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) - assert.NoError(t, err) - - err = db.AutoMigrate(&dbmodel.Segment{}, &dbmodel.SegmentMetadata{}) - assert.NoError(t, err) - - uniqueID := types.NewUniqueID() - collectionID := uniqueID.String() - testTopic := "test_topic" - segment := &dbmodel.Segment{ - ID: uniqueID.String(), - CollectionID: &collectionID, - Type: "test_type", - Scope: "test_scope", - Topic: &testTopic, - } - err = db.Create(segment).Error - assert.NoError(t, err) - - testKey := "test" - testValue := "test" - metadata := &dbmodel.SegmentMetadata{ - SegmentID: segment.ID, - Key: &testKey, - StrValue: &testValue, - } - err = db.Create(metadata).Error - assert.NoError(t, err) - - segmentDb := &segmentDb{ - db: db, - } - - // Test when all parameters are nil - segments, err := segmentDb.GetSegments(types.NilUniqueID(), nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Len(t, segments, 1) - assert.Equal(t, segment.ID, segments[0].Segment.ID) - assert.Equal(t, segment.CollectionID, segments[0].Segment.CollectionID) - assert.Equal(t, segment.Type, segments[0].Segment.Type) - assert.Equal(t, segment.Scope, segments[0].Segment.Scope) - assert.Equal(t, segment.Topic, segments[0].Segment.Topic) - assert.Len(t, segments[0].SegmentMetadata, 1) - assert.Equal(t, metadata.Key, segments[0].SegmentMetadata[0].Key) - assert.Equal(t, metadata.StrValue, segments[0].SegmentMetadata[0].StrValue) - - // Test when filtering by ID - segments, err = segmentDb.GetSegments(types.MustParse(segment.ID), nil, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Len(t, segments, 1) - assert.Equal(t, segment.ID, segments[0].Segment.ID) - - // Test when filtering by type - segments, err = segmentDb.GetSegments(types.NilUniqueID(), &segment.Type, nil, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Len(t, segments, 1) - assert.Equal(t, segment.ID, segments[0].Segment.ID) - - // Test when filtering by scope - segments, err = segmentDb.GetSegments(types.NilUniqueID(), nil, &segment.Scope, nil, types.NilUniqueID()) - assert.NoError(t, err) - assert.Len(t, segments, 1) - assert.Equal(t, segment.ID, segments[0].Segment.ID) - - // Test when filtering by topic - segments, err = segmentDb.GetSegments(types.NilUniqueID(), nil, nil, segment.Topic, types.NilUniqueID()) - assert.NoError(t, err) - assert.Len(t, segments, 1) - assert.Equal(t, segment.ID, segments[0].Segment.ID) - - // Test when filtering by collection ID - segments, err = segmentDb.GetSegments(types.NilUniqueID(), nil, nil, nil, types.MustParse(*segment.CollectionID)) - assert.NoError(t, err) - assert.Len(t, segments, 1) - assert.Equal(t, segment.ID, segments[0].Segment.ID) -} diff --git a/go/coordinator/internal/metastore/db/dao/tenant.go b/go/coordinator/internal/metastore/db/dao/tenant.go deleted file mode 100644 index 3fe759082ec..00000000000 --- a/go/coordinator/internal/metastore/db/dao/tenant.go +++ /dev/null @@ -1,38 +0,0 @@ -package dao - -import ( - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "gorm.io/gorm" -) - -type tenantDb struct { - db *gorm.DB -} - -var _ dbmodel.ITenantDb = &tenantDb{} - -func (s *tenantDb) DeleteAll() error { - return s.db.Where("1 = 1").Delete(&dbmodel.Tenant{}).Error -} - -func (s *tenantDb) GetAllTenants() ([]*dbmodel.Tenant, error) { - var tenants []*dbmodel.Tenant - - if err := s.db.Find(&tenants).Error; err != nil { - return nil, err - } - return tenants, nil -} - -func (s *tenantDb) GetTenants(tenantID string) ([]*dbmodel.Tenant, error) { - var tenants []*dbmodel.Tenant - - if err := s.db.Where("id = ?", tenantID).Find(&tenants).Error; err != nil { - return nil, err - } - return tenants, nil -} - -func (s *tenantDb) Insert(tenant *dbmodel.Tenant) error { - return s.db.Create(tenant).Error -} diff --git a/go/coordinator/internal/metastore/db/dbcore/core.go b/go/coordinator/internal/metastore/db/dbcore/core.go deleted file mode 100644 index 95d2885dfc4..00000000000 --- a/go/coordinator/internal/metastore/db/dbcore/core.go +++ /dev/null @@ -1,158 +0,0 @@ -package dbcore - -import ( - "context" - "fmt" - "reflect" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/types" - "github.com/pingcap/log" - "go.uber.org/zap" - "gorm.io/driver/postgres" - "gorm.io/driver/sqlite" - "gorm.io/gorm" - "gorm.io/gorm/logger" -) - -var ( - globalDB *gorm.DB -) - -type DBConfig struct { - Username string - Password string - Address string - Port int - DBName string - MaxIdleConns int - MaxOpenConns int -} - -func Connect(cfg DBConfig) (*gorm.DB, error) { - dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d sslmode=require", - cfg.Address, cfg.Username, cfg.Password, cfg.DBName, cfg.Port) - - ormLogger := logger.Default - ormLogger.LogMode(logger.Info) - db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{ - Logger: ormLogger, - CreateBatchSize: 100, - }) - if err != nil { - log.Error("fail to connect db", - zap.String("host", cfg.Address), - zap.String("database", cfg.DBName), - zap.Error(err)) - return nil, err - } - - idb, err := db.DB() - if err != nil { - log.Error("fail to create db instance", - zap.String("host", cfg.Address), - zap.String("database", cfg.DBName), - zap.Error(err)) - return nil, err - } - idb.SetMaxIdleConns(cfg.MaxIdleConns) - idb.SetMaxOpenConns(cfg.MaxOpenConns) - - globalDB = db - - log.Info("db connected success", - zap.String("host", cfg.Address), - zap.String("database", cfg.DBName), - zap.Error(err)) - - return db, nil -} - -// SetGlobalDB Only for test -func SetGlobalDB(db *gorm.DB) { - globalDB = db -} - -type ctxTransactionKey struct{} - -func CtxWithTransaction(ctx context.Context, tx *gorm.DB) context.Context { - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, ctxTransactionKey{}, tx) -} - -type txImpl struct{} - -func NewTxImpl() *txImpl { - return &txImpl{} -} - -func (*txImpl) Transaction(ctx context.Context, fn func(txctx context.Context) error) error { - db := globalDB.WithContext(ctx) - - return db.Transaction(func(tx *gorm.DB) error { - txCtx := CtxWithTransaction(ctx, tx) - return fn(txCtx) - }) -} - -func GetDB(ctx context.Context) *gorm.DB { - iface := ctx.Value(ctxTransactionKey{}) - - if iface != nil { - tx, ok := iface.(*gorm.DB) - if !ok { - log.Error("unexpect context value type", zap.Any("type", reflect.TypeOf(tx))) - return nil - } - - return tx - } - - return globalDB.WithContext(ctx) -} - -func ConfigDatabaseForTesting() *gorm.DB { - db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ - Logger: logger.Default.LogMode(logger.Info), - }) - if err != nil { - panic("failed to connect database") - } - SetGlobalDB(db) - // Setup tenant related tables - db.Migrator().DropTable(&dbmodel.Tenant{}) - db.Migrator().CreateTable(&dbmodel.Tenant{}) - db.Model(&dbmodel.Tenant{}).Create(&dbmodel.Tenant{ - ID: common.DefaultTenant, - }) - - // Setup database related tables - db.Migrator().DropTable(&dbmodel.Database{}) - db.Migrator().CreateTable(&dbmodel.Database{}) - - db.Model(&dbmodel.Database{}).Create(&dbmodel.Database{ - ID: types.NilUniqueID().String(), - Name: common.DefaultDatabase, - TenantID: common.DefaultTenant, - }) - - // Setup collection related tables - db.Migrator().DropTable(&dbmodel.Collection{}) - db.Migrator().DropTable(&dbmodel.CollectionMetadata{}) - db.Migrator().CreateTable(&dbmodel.Collection{}) - db.Migrator().CreateTable(&dbmodel.CollectionMetadata{}) - - // Setup segment related tables - db.Migrator().DropTable(&dbmodel.Segment{}) - db.Migrator().DropTable(&dbmodel.SegmentMetadata{}) - db.Migrator().CreateTable(&dbmodel.Segment{}) - db.Migrator().CreateTable(&dbmodel.SegmentMetadata{}) - - // Setup notification related tables - db.Migrator().DropTable(&dbmodel.Notification{}) - db.Migrator().CreateTable(&dbmodel.Notification{}) - return db -} diff --git a/go/coordinator/internal/metastore/db/dbmodel/collection.go b/go/coordinator/internal/metastore/db/dbmodel/collection.go deleted file mode 100644 index 46f00474d4e..00000000000 --- a/go/coordinator/internal/metastore/db/dbmodel/collection.go +++ /dev/null @@ -1,39 +0,0 @@ -package dbmodel - -import ( - "time" - - "github.com/chroma/chroma-coordinator/internal/types" -) - -type Collection struct { - ID string `gorm:"id;primaryKey"` - Name *string `gorm:"name;unique"` - Topic *string `gorm:"topic"` - Dimension *int32 `gorm:"dimension"` - DatabaseID string `gorm:"database_id"` - Ts types.Timestamp `gorm:"ts;type:bigint;default:0"` - IsDeleted bool `gorm:"is_deleted;type:bool;default:false"` - CreatedAt time.Time `gorm:"created_at;type:timestamp;not null;default:current_timestamp"` - UpdatedAt time.Time `gorm:"updated_at;type:timestamp;not null;default:current_timestamp"` -} - -func (v Collection) TableName() string { - return "collections" -} - -type CollectionAndMetadata struct { - Collection *Collection - CollectionMetadata []*CollectionMetadata - TenantID string - DatabaseName string -} - -//go:generate mockery --name=ICollectionDb -type ICollectionDb interface { - GetCollections(collectionID *string, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*CollectionAndMetadata, error) - DeleteCollectionByID(collectionID string) error - Insert(in *Collection) error - Update(in *Collection) error - DeleteAll() error -} diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/ICollectionDb.go b/go/coordinator/internal/metastore/db/dbmodel/mocks/ICollectionDb.go deleted file mode 100644 index 109db42818b..00000000000 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/ICollectionDb.go +++ /dev/null @@ -1,109 +0,0 @@ -// Code generated by mockery v2.33.3. DO NOT EDIT. - -package mocks - -import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - mock "github.com/stretchr/testify/mock" -) - -// ICollectionDb is an autogenerated mock type for the ICollectionDb type -type ICollectionDb struct { - mock.Mock -} - -// DeleteAll provides a mock function with given fields: -func (_m *ICollectionDb) DeleteAll() error { - ret := _m.Called() - - var r0 error - if rf, ok := ret.Get(0).(func() error); ok { - r0 = rf() - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// DeleteCollectionByID provides a mock function with given fields: collectionID -func (_m *ICollectionDb) DeleteCollectionByID(collectionID string) error { - ret := _m.Called(collectionID) - - var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(collectionID) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// GetCollections provides a mock function with given fields: collectionID, collectionName, collectionTopic, tenantID, databaseName -func (_m *ICollectionDb) GetCollections(collectionID *string, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*dbmodel.CollectionAndMetadata, error) { - ret := _m.Called(collectionID, collectionName, collectionTopic, tenantID, databaseName) - - var r0 []*dbmodel.CollectionAndMetadata - var r1 error - if rf, ok := ret.Get(0).(func(*string, *string, *string, string, string) ([]*dbmodel.CollectionAndMetadata, error)); ok { - return rf(collectionID, collectionName, collectionTopic, tenantID, databaseName) - } - if rf, ok := ret.Get(0).(func(*string, *string, *string, string, string) []*dbmodel.CollectionAndMetadata); ok { - r0 = rf(collectionID, collectionName, collectionTopic, tenantID, databaseName) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).([]*dbmodel.CollectionAndMetadata) - } - } - - if rf, ok := ret.Get(1).(func(*string, *string, *string, string, string) error); ok { - r1 = rf(collectionID, collectionName, collectionTopic, tenantID, databaseName) - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// Insert provides a mock function with given fields: in -func (_m *ICollectionDb) Insert(in *dbmodel.Collection) error { - ret := _m.Called(in) - - var r0 error - if rf, ok := ret.Get(0).(func(*dbmodel.Collection) error); ok { - r0 = rf(in) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// Update provides a mock function with given fields: in -func (_m *ICollectionDb) Update(in *dbmodel.Collection) error { - ret := _m.Called(in) - - var r0 error - if rf, ok := ret.Get(0).(func(*dbmodel.Collection) error); ok { - r0 = rf(in) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// NewICollectionDb creates a new instance of ICollectionDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -// The first argument is typically a *testing.T value. -func NewICollectionDb(t interface { - mock.TestingT - Cleanup(func()) -}) *ICollectionDb { - mock := &ICollectionDb{} - mock.Mock.Test(t) - - t.Cleanup(func() { mock.AssertExpectations(t) }) - - return mock -} diff --git a/go/coordinator/internal/metastore/db/dbmodel/segment.go b/go/coordinator/internal/metastore/db/dbmodel/segment.go deleted file mode 100644 index 0967436e11e..00000000000 --- a/go/coordinator/internal/metastore/db/dbmodel/segment.go +++ /dev/null @@ -1,45 +0,0 @@ -package dbmodel - -import ( - "time" - - "github.com/chroma/chroma-coordinator/internal/types" -) - -type Segment struct { - ID string `gorm:"id;primaryKey"` - Type string `gorm:"type;type:string;not null"` - Scope string `gorm:"scope"` - Topic *string `gorm:"topic"` - Ts types.Timestamp `gorm:"ts;type:bigint;default:0"` - IsDeleted bool `gorm:"is_deleted;type:bool;default:false"` - CreatedAt time.Time `gorm:"created_at;type:timestamp;not null;default:current_timestamp"` - UpdatedAt time.Time `gorm:"updated_at;type:timestamp;not null;default:current_timestamp"` - CollectionID *string `gorm:"collection_id"` -} - -func (s Segment) TableName() string { - return "segments" -} - -type SegmentAndMetadata struct { - Segment *Segment - SegmentMetadata []*SegmentMetadata -} - -type UpdateSegment struct { - ID string - Topic *string - ResetTopic bool - Collection *string - ResetCollection bool -} - -//go:generate mockery --name=ISegmentDb -type ISegmentDb interface { - GetSegments(id types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*SegmentAndMetadata, error) - DeleteSegmentByID(id string) error - Insert(*Segment) error - Update(*UpdateSegment) error - DeleteAll() error -} diff --git a/go/coordinator/internal/metastore/db/dbmodel/tenant.go b/go/coordinator/internal/metastore/db/dbmodel/tenant.go deleted file mode 100644 index bb15ed5153e..00000000000 --- a/go/coordinator/internal/metastore/db/dbmodel/tenant.go +++ /dev/null @@ -1,27 +0,0 @@ -package dbmodel - -import ( - "time" - - "github.com/chroma/chroma-coordinator/internal/types" -) - -type Tenant struct { - ID string `gorm:"id;primaryKey;unique"` - Ts types.Timestamp `gorm:"ts;type:bigint;default:0"` - IsDeleted bool `gorm:"is_deleted;type:bool;default:false"` - CreatedAt time.Time `gorm:"created_at;type:timestamp;not null;default:current_timestamp"` - UpdatedAt time.Time `gorm:"updated_at;type:timestamp;not null;default:current_timestamp"` -} - -func (v Tenant) TableName() string { - return "tenants" -} - -//go:generate mockery --name=ITenantDb -type ITenantDb interface { - GetAllTenants() ([]*Tenant, error) - GetTenants(tenantID string) ([]*Tenant, error) - Insert(in *Tenant) error - DeleteAll() error -} diff --git a/go/coordinator/internal/proto/coordinatorpb/coordinator.pb.go b/go/coordinator/internal/proto/coordinatorpb/coordinator.pb.go deleted file mode 100644 index be93392c304..00000000000 --- a/go/coordinator/internal/proto/coordinatorpb/coordinator.pb.go +++ /dev/null @@ -1,1865 +0,0 @@ -// Code generated by protoc-gen-go. DO NOT EDIT. -// versions: -// protoc-gen-go v1.31.0 -// protoc v4.23.4 -// source: chromadb/proto/coordinator.proto - -package coordinatorpb - -import ( - protoreflect "google.golang.org/protobuf/reflect/protoreflect" - protoimpl "google.golang.org/protobuf/runtime/protoimpl" - emptypb "google.golang.org/protobuf/types/known/emptypb" - reflect "reflect" - sync "sync" -) - -const ( - // Verify that this generated code is sufficiently up-to-date. - _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) - // Verify that runtime/protoimpl is sufficiently up-to-date. - _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) -) - -type CreateDatabaseRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` - Tenant string `protobuf:"bytes,3,opt,name=tenant,proto3" json:"tenant,omitempty"` -} - -func (x *CreateDatabaseRequest) Reset() { - *x = CreateDatabaseRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[0] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *CreateDatabaseRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*CreateDatabaseRequest) ProtoMessage() {} - -func (x *CreateDatabaseRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[0] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use CreateDatabaseRequest.ProtoReflect.Descriptor instead. -func (*CreateDatabaseRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{0} -} - -func (x *CreateDatabaseRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (x *CreateDatabaseRequest) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -func (x *CreateDatabaseRequest) GetTenant() string { - if x != nil { - return x.Tenant - } - return "" -} - -type GetDatabaseRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` - Tenant string `protobuf:"bytes,2,opt,name=tenant,proto3" json:"tenant,omitempty"` -} - -func (x *GetDatabaseRequest) Reset() { - *x = GetDatabaseRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[1] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetDatabaseRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetDatabaseRequest) ProtoMessage() {} - -func (x *GetDatabaseRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[1] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetDatabaseRequest.ProtoReflect.Descriptor instead. -func (*GetDatabaseRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{1} -} - -func (x *GetDatabaseRequest) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -func (x *GetDatabaseRequest) GetTenant() string { - if x != nil { - return x.Tenant - } - return "" -} - -type GetDatabaseResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Database *Database `protobuf:"bytes,1,opt,name=database,proto3" json:"database,omitempty"` - Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *GetDatabaseResponse) Reset() { - *x = GetDatabaseResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[2] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetDatabaseResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetDatabaseResponse) ProtoMessage() {} - -func (x *GetDatabaseResponse) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[2] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetDatabaseResponse.ProtoReflect.Descriptor instead. -func (*GetDatabaseResponse) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{2} -} - -func (x *GetDatabaseResponse) GetDatabase() *Database { - if x != nil { - return x.Database - } - return nil -} - -func (x *GetDatabaseResponse) GetStatus() *Status { - if x != nil { - return x.Status - } - return nil -} - -type CreateTenantRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` // Names are globally unique -} - -func (x *CreateTenantRequest) Reset() { - *x = CreateTenantRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[3] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *CreateTenantRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*CreateTenantRequest) ProtoMessage() {} - -func (x *CreateTenantRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[3] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use CreateTenantRequest.ProtoReflect.Descriptor instead. -func (*CreateTenantRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{3} -} - -func (x *CreateTenantRequest) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -type GetTenantRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` -} - -func (x *GetTenantRequest) Reset() { - *x = GetTenantRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[4] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetTenantRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetTenantRequest) ProtoMessage() {} - -func (x *GetTenantRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[4] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetTenantRequest.ProtoReflect.Descriptor instead. -func (*GetTenantRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{4} -} - -func (x *GetTenantRequest) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -type GetTenantResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Tenant *Tenant `protobuf:"bytes,1,opt,name=tenant,proto3" json:"tenant,omitempty"` - Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *GetTenantResponse) Reset() { - *x = GetTenantResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[5] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetTenantResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetTenantResponse) ProtoMessage() {} - -func (x *GetTenantResponse) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[5] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetTenantResponse.ProtoReflect.Descriptor instead. -func (*GetTenantResponse) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{5} -} - -func (x *GetTenantResponse) GetTenant() *Tenant { - if x != nil { - return x.Tenant - } - return nil -} - -func (x *GetTenantResponse) GetStatus() *Status { - if x != nil { - return x.Status - } - return nil -} - -type CreateSegmentRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Segment *Segment `protobuf:"bytes,1,opt,name=segment,proto3" json:"segment,omitempty"` -} - -func (x *CreateSegmentRequest) Reset() { - *x = CreateSegmentRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[6] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *CreateSegmentRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*CreateSegmentRequest) ProtoMessage() {} - -func (x *CreateSegmentRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[6] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use CreateSegmentRequest.ProtoReflect.Descriptor instead. -func (*CreateSegmentRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{6} -} - -func (x *CreateSegmentRequest) GetSegment() *Segment { - if x != nil { - return x.Segment - } - return nil -} - -type DeleteSegmentRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` -} - -func (x *DeleteSegmentRequest) Reset() { - *x = DeleteSegmentRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[7] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *DeleteSegmentRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*DeleteSegmentRequest) ProtoMessage() {} - -func (x *DeleteSegmentRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[7] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use DeleteSegmentRequest.ProtoReflect.Descriptor instead. -func (*DeleteSegmentRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{7} -} - -func (x *DeleteSegmentRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -type GetSegmentsRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id *string `protobuf:"bytes,1,opt,name=id,proto3,oneof" json:"id,omitempty"` - Type *string `protobuf:"bytes,2,opt,name=type,proto3,oneof" json:"type,omitempty"` - Scope *SegmentScope `protobuf:"varint,3,opt,name=scope,proto3,enum=chroma.SegmentScope,oneof" json:"scope,omitempty"` - Topic *string `protobuf:"bytes,4,opt,name=topic,proto3,oneof" json:"topic,omitempty"` - Collection *string `protobuf:"bytes,5,opt,name=collection,proto3,oneof" json:"collection,omitempty"` // Collection ID -} - -func (x *GetSegmentsRequest) Reset() { - *x = GetSegmentsRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[8] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetSegmentsRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetSegmentsRequest) ProtoMessage() {} - -func (x *GetSegmentsRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[8] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetSegmentsRequest.ProtoReflect.Descriptor instead. -func (*GetSegmentsRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{8} -} - -func (x *GetSegmentsRequest) GetId() string { - if x != nil && x.Id != nil { - return *x.Id - } - return "" -} - -func (x *GetSegmentsRequest) GetType() string { - if x != nil && x.Type != nil { - return *x.Type - } - return "" -} - -func (x *GetSegmentsRequest) GetScope() SegmentScope { - if x != nil && x.Scope != nil { - return *x.Scope - } - return SegmentScope_VECTOR -} - -func (x *GetSegmentsRequest) GetTopic() string { - if x != nil && x.Topic != nil { - return *x.Topic - } - return "" -} - -func (x *GetSegmentsRequest) GetCollection() string { - if x != nil && x.Collection != nil { - return *x.Collection - } - return "" -} - -type GetSegmentsResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Segments []*Segment `protobuf:"bytes,1,rep,name=segments,proto3" json:"segments,omitempty"` - Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *GetSegmentsResponse) Reset() { - *x = GetSegmentsResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[9] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetSegmentsResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetSegmentsResponse) ProtoMessage() {} - -func (x *GetSegmentsResponse) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[9] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetSegmentsResponse.ProtoReflect.Descriptor instead. -func (*GetSegmentsResponse) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{9} -} - -func (x *GetSegmentsResponse) GetSegments() []*Segment { - if x != nil { - return x.Segments - } - return nil -} - -func (x *GetSegmentsResponse) GetStatus() *Status { - if x != nil { - return x.Status - } - return nil -} - -type UpdateSegmentRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - // Types that are assignable to TopicUpdate: - // - // *UpdateSegmentRequest_Topic - // *UpdateSegmentRequest_ResetTopic - TopicUpdate isUpdateSegmentRequest_TopicUpdate `protobuf_oneof:"topic_update"` - // Types that are assignable to CollectionUpdate: - // - // *UpdateSegmentRequest_Collection - // *UpdateSegmentRequest_ResetCollection - CollectionUpdate isUpdateSegmentRequest_CollectionUpdate `protobuf_oneof:"collection_update"` - // Types that are assignable to MetadataUpdate: - // - // *UpdateSegmentRequest_Metadata - // *UpdateSegmentRequest_ResetMetadata - MetadataUpdate isUpdateSegmentRequest_MetadataUpdate `protobuf_oneof:"metadata_update"` -} - -func (x *UpdateSegmentRequest) Reset() { - *x = UpdateSegmentRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[10] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *UpdateSegmentRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*UpdateSegmentRequest) ProtoMessage() {} - -func (x *UpdateSegmentRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[10] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use UpdateSegmentRequest.ProtoReflect.Descriptor instead. -func (*UpdateSegmentRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{10} -} - -func (x *UpdateSegmentRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (m *UpdateSegmentRequest) GetTopicUpdate() isUpdateSegmentRequest_TopicUpdate { - if m != nil { - return m.TopicUpdate - } - return nil -} - -func (x *UpdateSegmentRequest) GetTopic() string { - if x, ok := x.GetTopicUpdate().(*UpdateSegmentRequest_Topic); ok { - return x.Topic - } - return "" -} - -func (x *UpdateSegmentRequest) GetResetTopic() bool { - if x, ok := x.GetTopicUpdate().(*UpdateSegmentRequest_ResetTopic); ok { - return x.ResetTopic - } - return false -} - -func (m *UpdateSegmentRequest) GetCollectionUpdate() isUpdateSegmentRequest_CollectionUpdate { - if m != nil { - return m.CollectionUpdate - } - return nil -} - -func (x *UpdateSegmentRequest) GetCollection() string { - if x, ok := x.GetCollectionUpdate().(*UpdateSegmentRequest_Collection); ok { - return x.Collection - } - return "" -} - -func (x *UpdateSegmentRequest) GetResetCollection() bool { - if x, ok := x.GetCollectionUpdate().(*UpdateSegmentRequest_ResetCollection); ok { - return x.ResetCollection - } - return false -} - -func (m *UpdateSegmentRequest) GetMetadataUpdate() isUpdateSegmentRequest_MetadataUpdate { - if m != nil { - return m.MetadataUpdate - } - return nil -} - -func (x *UpdateSegmentRequest) GetMetadata() *UpdateMetadata { - if x, ok := x.GetMetadataUpdate().(*UpdateSegmentRequest_Metadata); ok { - return x.Metadata - } - return nil -} - -func (x *UpdateSegmentRequest) GetResetMetadata() bool { - if x, ok := x.GetMetadataUpdate().(*UpdateSegmentRequest_ResetMetadata); ok { - return x.ResetMetadata - } - return false -} - -type isUpdateSegmentRequest_TopicUpdate interface { - isUpdateSegmentRequest_TopicUpdate() -} - -type UpdateSegmentRequest_Topic struct { - Topic string `protobuf:"bytes,2,opt,name=topic,proto3,oneof"` -} - -type UpdateSegmentRequest_ResetTopic struct { - ResetTopic bool `protobuf:"varint,3,opt,name=reset_topic,json=resetTopic,proto3,oneof"` -} - -func (*UpdateSegmentRequest_Topic) isUpdateSegmentRequest_TopicUpdate() {} - -func (*UpdateSegmentRequest_ResetTopic) isUpdateSegmentRequest_TopicUpdate() {} - -type isUpdateSegmentRequest_CollectionUpdate interface { - isUpdateSegmentRequest_CollectionUpdate() -} - -type UpdateSegmentRequest_Collection struct { - Collection string `protobuf:"bytes,4,opt,name=collection,proto3,oneof"` -} - -type UpdateSegmentRequest_ResetCollection struct { - ResetCollection bool `protobuf:"varint,5,opt,name=reset_collection,json=resetCollection,proto3,oneof"` -} - -func (*UpdateSegmentRequest_Collection) isUpdateSegmentRequest_CollectionUpdate() {} - -func (*UpdateSegmentRequest_ResetCollection) isUpdateSegmentRequest_CollectionUpdate() {} - -type isUpdateSegmentRequest_MetadataUpdate interface { - isUpdateSegmentRequest_MetadataUpdate() -} - -type UpdateSegmentRequest_Metadata struct { - Metadata *UpdateMetadata `protobuf:"bytes,6,opt,name=metadata,proto3,oneof"` -} - -type UpdateSegmentRequest_ResetMetadata struct { - ResetMetadata bool `protobuf:"varint,7,opt,name=reset_metadata,json=resetMetadata,proto3,oneof"` -} - -func (*UpdateSegmentRequest_Metadata) isUpdateSegmentRequest_MetadataUpdate() {} - -func (*UpdateSegmentRequest_ResetMetadata) isUpdateSegmentRequest_MetadataUpdate() {} - -type CreateCollectionRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` - Metadata *UpdateMetadata `protobuf:"bytes,3,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` - Dimension *int32 `protobuf:"varint,4,opt,name=dimension,proto3,oneof" json:"dimension,omitempty"` - GetOrCreate *bool `protobuf:"varint,5,opt,name=get_or_create,json=getOrCreate,proto3,oneof" json:"get_or_create,omitempty"` - Tenant string `protobuf:"bytes,6,opt,name=tenant,proto3" json:"tenant,omitempty"` - Database string `protobuf:"bytes,7,opt,name=database,proto3" json:"database,omitempty"` -} - -func (x *CreateCollectionRequest) Reset() { - *x = CreateCollectionRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[11] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *CreateCollectionRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*CreateCollectionRequest) ProtoMessage() {} - -func (x *CreateCollectionRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[11] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use CreateCollectionRequest.ProtoReflect.Descriptor instead. -func (*CreateCollectionRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{11} -} - -func (x *CreateCollectionRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (x *CreateCollectionRequest) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -func (x *CreateCollectionRequest) GetMetadata() *UpdateMetadata { - if x != nil { - return x.Metadata - } - return nil -} - -func (x *CreateCollectionRequest) GetDimension() int32 { - if x != nil && x.Dimension != nil { - return *x.Dimension - } - return 0 -} - -func (x *CreateCollectionRequest) GetGetOrCreate() bool { - if x != nil && x.GetOrCreate != nil { - return *x.GetOrCreate - } - return false -} - -func (x *CreateCollectionRequest) GetTenant() string { - if x != nil { - return x.Tenant - } - return "" -} - -func (x *CreateCollectionRequest) GetDatabase() string { - if x != nil { - return x.Database - } - return "" -} - -type CreateCollectionResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Collection *Collection `protobuf:"bytes,1,opt,name=collection,proto3" json:"collection,omitempty"` - Created bool `protobuf:"varint,2,opt,name=created,proto3" json:"created,omitempty"` - Status *Status `protobuf:"bytes,3,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *CreateCollectionResponse) Reset() { - *x = CreateCollectionResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[12] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *CreateCollectionResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*CreateCollectionResponse) ProtoMessage() {} - -func (x *CreateCollectionResponse) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[12] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use CreateCollectionResponse.ProtoReflect.Descriptor instead. -func (*CreateCollectionResponse) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{12} -} - -func (x *CreateCollectionResponse) GetCollection() *Collection { - if x != nil { - return x.Collection - } - return nil -} - -func (x *CreateCollectionResponse) GetCreated() bool { - if x != nil { - return x.Created - } - return false -} - -func (x *CreateCollectionResponse) GetStatus() *Status { - if x != nil { - return x.Status - } - return nil -} - -type DeleteCollectionRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Tenant string `protobuf:"bytes,2,opt,name=tenant,proto3" json:"tenant,omitempty"` - Database string `protobuf:"bytes,3,opt,name=database,proto3" json:"database,omitempty"` -} - -func (x *DeleteCollectionRequest) Reset() { - *x = DeleteCollectionRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[13] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *DeleteCollectionRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*DeleteCollectionRequest) ProtoMessage() {} - -func (x *DeleteCollectionRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[13] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use DeleteCollectionRequest.ProtoReflect.Descriptor instead. -func (*DeleteCollectionRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{13} -} - -func (x *DeleteCollectionRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (x *DeleteCollectionRequest) GetTenant() string { - if x != nil { - return x.Tenant - } - return "" -} - -func (x *DeleteCollectionRequest) GetDatabase() string { - if x != nil { - return x.Database - } - return "" -} - -type GetCollectionsRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id *string `protobuf:"bytes,1,opt,name=id,proto3,oneof" json:"id,omitempty"` - Name *string `protobuf:"bytes,2,opt,name=name,proto3,oneof" json:"name,omitempty"` - Topic *string `protobuf:"bytes,3,opt,name=topic,proto3,oneof" json:"topic,omitempty"` - Tenant string `protobuf:"bytes,4,opt,name=tenant,proto3" json:"tenant,omitempty"` - Database string `protobuf:"bytes,5,opt,name=database,proto3" json:"database,omitempty"` -} - -func (x *GetCollectionsRequest) Reset() { - *x = GetCollectionsRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[14] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetCollectionsRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetCollectionsRequest) ProtoMessage() {} - -func (x *GetCollectionsRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[14] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetCollectionsRequest.ProtoReflect.Descriptor instead. -func (*GetCollectionsRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{14} -} - -func (x *GetCollectionsRequest) GetId() string { - if x != nil && x.Id != nil { - return *x.Id - } - return "" -} - -func (x *GetCollectionsRequest) GetName() string { - if x != nil && x.Name != nil { - return *x.Name - } - return "" -} - -func (x *GetCollectionsRequest) GetTopic() string { - if x != nil && x.Topic != nil { - return *x.Topic - } - return "" -} - -func (x *GetCollectionsRequest) GetTenant() string { - if x != nil { - return x.Tenant - } - return "" -} - -func (x *GetCollectionsRequest) GetDatabase() string { - if x != nil { - return x.Database - } - return "" -} - -type GetCollectionsResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Collections []*Collection `protobuf:"bytes,1,rep,name=collections,proto3" json:"collections,omitempty"` - Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *GetCollectionsResponse) Reset() { - *x = GetCollectionsResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[15] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *GetCollectionsResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetCollectionsResponse) ProtoMessage() {} - -func (x *GetCollectionsResponse) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[15] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetCollectionsResponse.ProtoReflect.Descriptor instead. -func (*GetCollectionsResponse) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{15} -} - -func (x *GetCollectionsResponse) GetCollections() []*Collection { - if x != nil { - return x.Collections - } - return nil -} - -func (x *GetCollectionsResponse) GetStatus() *Status { - if x != nil { - return x.Status - } - return nil -} - -type UpdateCollectionRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Topic *string `protobuf:"bytes,2,opt,name=topic,proto3,oneof" json:"topic,omitempty"` - Name *string `protobuf:"bytes,3,opt,name=name,proto3,oneof" json:"name,omitempty"` - Dimension *int32 `protobuf:"varint,4,opt,name=dimension,proto3,oneof" json:"dimension,omitempty"` - // Types that are assignable to MetadataUpdate: - // - // *UpdateCollectionRequest_Metadata - // *UpdateCollectionRequest_ResetMetadata - MetadataUpdate isUpdateCollectionRequest_MetadataUpdate `protobuf_oneof:"metadata_update"` -} - -func (x *UpdateCollectionRequest) Reset() { - *x = UpdateCollectionRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[16] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *UpdateCollectionRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*UpdateCollectionRequest) ProtoMessage() {} - -func (x *UpdateCollectionRequest) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[16] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use UpdateCollectionRequest.ProtoReflect.Descriptor instead. -func (*UpdateCollectionRequest) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{16} -} - -func (x *UpdateCollectionRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (x *UpdateCollectionRequest) GetTopic() string { - if x != nil && x.Topic != nil { - return *x.Topic - } - return "" -} - -func (x *UpdateCollectionRequest) GetName() string { - if x != nil && x.Name != nil { - return *x.Name - } - return "" -} - -func (x *UpdateCollectionRequest) GetDimension() int32 { - if x != nil && x.Dimension != nil { - return *x.Dimension - } - return 0 -} - -func (m *UpdateCollectionRequest) GetMetadataUpdate() isUpdateCollectionRequest_MetadataUpdate { - if m != nil { - return m.MetadataUpdate - } - return nil -} - -func (x *UpdateCollectionRequest) GetMetadata() *UpdateMetadata { - if x, ok := x.GetMetadataUpdate().(*UpdateCollectionRequest_Metadata); ok { - return x.Metadata - } - return nil -} - -func (x *UpdateCollectionRequest) GetResetMetadata() bool { - if x, ok := x.GetMetadataUpdate().(*UpdateCollectionRequest_ResetMetadata); ok { - return x.ResetMetadata - } - return false -} - -type isUpdateCollectionRequest_MetadataUpdate interface { - isUpdateCollectionRequest_MetadataUpdate() -} - -type UpdateCollectionRequest_Metadata struct { - Metadata *UpdateMetadata `protobuf:"bytes,5,opt,name=metadata,proto3,oneof"` -} - -type UpdateCollectionRequest_ResetMetadata struct { - ResetMetadata bool `protobuf:"varint,6,opt,name=reset_metadata,json=resetMetadata,proto3,oneof"` -} - -func (*UpdateCollectionRequest_Metadata) isUpdateCollectionRequest_MetadataUpdate() {} - -func (*UpdateCollectionRequest_ResetMetadata) isUpdateCollectionRequest_MetadataUpdate() {} - -type Notification struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id int64 `protobuf:"varint,1,opt,name=id,proto3" json:"id,omitempty"` - CollectionId string `protobuf:"bytes,2,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` - Type string `protobuf:"bytes,3,opt,name=type,proto3" json:"type,omitempty"` - Status string `protobuf:"bytes,4,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *Notification) Reset() { - *x = Notification{} - if protoimpl.UnsafeEnabled { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[17] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *Notification) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*Notification) ProtoMessage() {} - -func (x *Notification) ProtoReflect() protoreflect.Message { - mi := &file_chromadb_proto_coordinator_proto_msgTypes[17] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use Notification.ProtoReflect.Descriptor instead. -func (*Notification) Descriptor() ([]byte, []int) { - return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{17} -} - -func (x *Notification) GetId() int64 { - if x != nil { - return x.Id - } - return 0 -} - -func (x *Notification) GetCollectionId() string { - if x != nil { - return x.CollectionId - } - return "" -} - -func (x *Notification) GetType() string { - if x != nil { - return x.Type - } - return "" -} - -func (x *Notification) GetStatus() string { - if x != nil { - return x.Status - } - return "" -} - -var File_chromadb_proto_coordinator_proto protoreflect.FileDescriptor - -var file_chromadb_proto_coordinator_proto_rawDesc = []byte{ - 0x0a, 0x20, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x64, 0x62, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x12, 0x06, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x1a, 0x1b, 0x63, 0x68, 0x72, 0x6f, - 0x6d, 0x61, 0x64, 0x62, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x1b, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x53, 0x0a, 0x15, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x44, 0x61, - 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, - 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, - 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, - 0x65, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x22, 0x40, 0x0a, 0x12, 0x47, 0x65, 0x74, - 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, - 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, - 0x61, 0x6d, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x22, 0x6b, 0x0a, 0x13, 0x47, - 0x65, 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, - 0x73, 0x65, 0x12, 0x2c, 0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x01, - 0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x44, 0x61, - 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, - 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, - 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, - 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x29, 0x0a, 0x13, 0x43, 0x72, 0x65, 0x61, - 0x74, 0x65, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, - 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, - 0x61, 0x6d, 0x65, 0x22, 0x26, 0x0a, 0x10, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, - 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x63, 0x0a, 0x11, 0x47, - 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, - 0x12, 0x26, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, - 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, - 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, - 0x75, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, - 0x22, 0x41, 0x0a, 0x14, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, - 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x29, 0x0a, 0x07, 0x73, 0x65, 0x67, 0x6d, - 0x65, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0f, 0x2e, 0x63, 0x68, 0x72, 0x6f, - 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x07, 0x73, 0x65, 0x67, 0x6d, - 0x65, 0x6e, 0x74, 0x22, 0x26, 0x0a, 0x14, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, 0x65, 0x67, - 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, - 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22, 0xe6, 0x01, 0x0a, 0x12, - 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, - 0x73, 0x74, 0x12, 0x13, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, - 0x52, 0x02, 0x69, 0x64, 0x88, 0x01, 0x01, 0x12, 0x17, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, - 0x02, 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x88, 0x01, 0x01, - 0x12, 0x2f, 0x0a, 0x05, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0e, 0x32, - 0x14, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, - 0x53, 0x63, 0x6f, 0x70, 0x65, 0x48, 0x02, 0x52, 0x05, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x88, 0x01, - 0x01, 0x12, 0x19, 0x0a, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, - 0x48, 0x03, 0x52, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x88, 0x01, 0x01, 0x12, 0x23, 0x0a, 0x0a, - 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, - 0x48, 0x04, 0x52, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x88, 0x01, - 0x01, 0x42, 0x05, 0x0a, 0x03, 0x5f, 0x69, 0x64, 0x42, 0x07, 0x0a, 0x05, 0x5f, 0x74, 0x79, 0x70, - 0x65, 0x42, 0x08, 0x0a, 0x06, 0x5f, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x42, 0x08, 0x0a, 0x06, 0x5f, - 0x74, 0x6f, 0x70, 0x69, 0x63, 0x42, 0x0d, 0x0a, 0x0b, 0x5f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, - 0x74, 0x69, 0x6f, 0x6e, 0x22, 0x6a, 0x0a, 0x13, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, - 0x6e, 0x74, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x2b, 0x0a, 0x08, 0x73, - 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x0f, 0x2e, - 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x08, - 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, - 0x75, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, - 0x22, 0xc7, 0x02, 0x0a, 0x14, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, - 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, - 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x16, 0x0a, 0x05, 0x74, 0x6f, 0x70, - 0x69, 0x63, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, 0x52, 0x05, 0x74, 0x6f, 0x70, 0x69, - 0x63, 0x12, 0x21, 0x0a, 0x0b, 0x72, 0x65, 0x73, 0x65, 0x74, 0x5f, 0x74, 0x6f, 0x70, 0x69, 0x63, - 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x48, 0x00, 0x52, 0x0a, 0x72, 0x65, 0x73, 0x65, 0x74, 0x54, - 0x6f, 0x70, 0x69, 0x63, 0x12, 0x20, 0x0a, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, - 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, - 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x2b, 0x0a, 0x10, 0x72, 0x65, 0x73, 0x65, 0x74, 0x5f, - 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x05, 0x20, 0x01, 0x28, 0x08, - 0x48, 0x01, 0x52, 0x0f, 0x72, 0x65, 0x73, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, - 0x69, 0x6f, 0x6e, 0x12, 0x34, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, - 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, - 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x02, 0x52, - 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x12, 0x27, 0x0a, 0x0e, 0x72, 0x65, 0x73, - 0x65, 0x74, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x07, 0x20, 0x01, 0x28, - 0x08, 0x48, 0x02, 0x52, 0x0d, 0x72, 0x65, 0x73, 0x65, 0x74, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, - 0x74, 0x61, 0x42, 0x0e, 0x0a, 0x0c, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x5f, 0x75, 0x70, 0x64, 0x61, - 0x74, 0x65, 0x42, 0x13, 0x0a, 0x11, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, - 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x42, 0x11, 0x0a, 0x0f, 0x6d, 0x65, 0x74, 0x61, 0x64, - 0x61, 0x74, 0x61, 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x22, 0xa3, 0x02, 0x0a, 0x17, 0x43, - 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, - 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x37, 0x0a, 0x08, 0x6d, 0x65, - 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, - 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, - 0x64, 0x61, 0x74, 0x61, 0x48, 0x00, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, - 0x88, 0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, - 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x48, 0x01, 0x52, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, - 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x27, 0x0a, 0x0d, 0x67, 0x65, 0x74, 0x5f, 0x6f, 0x72, - 0x5f, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x08, 0x48, 0x02, 0x52, - 0x0b, 0x67, 0x65, 0x74, 0x4f, 0x72, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x88, 0x01, 0x01, 0x12, - 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, - 0x61, 0x73, 0x65, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, - 0x61, 0x73, 0x65, 0x42, 0x0b, 0x0a, 0x09, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, - 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x42, 0x10, - 0x0a, 0x0e, 0x5f, 0x67, 0x65, 0x74, 0x5f, 0x6f, 0x72, 0x5f, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, - 0x22, 0x90, 0x01, 0x0a, 0x18, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, - 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x32, 0x0a, - 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, - 0x0b, 0x32, 0x12, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x6f, 0x6c, 0x6c, 0x65, - 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, - 0x6e, 0x12, 0x18, 0x0a, 0x07, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, 0x64, 0x18, 0x02, 0x20, 0x01, - 0x28, 0x08, 0x52, 0x07, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, 0x64, 0x12, 0x26, 0x0a, 0x06, 0x73, - 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, - 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, - 0x74, 0x75, 0x73, 0x22, 0x5d, 0x0a, 0x17, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x43, 0x6f, 0x6c, - 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, - 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x16, - 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, - 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, - 0x73, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, - 0x73, 0x65, 0x22, 0xae, 0x01, 0x0a, 0x15, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, - 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x13, 0x0a, 0x02, - 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, 0x52, 0x02, 0x69, 0x64, 0x88, 0x01, - 0x01, 0x12, 0x17, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x48, - 0x01, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x88, 0x01, 0x01, 0x12, 0x19, 0x0a, 0x05, 0x74, 0x6f, - 0x70, 0x69, 0x63, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x48, 0x02, 0x52, 0x05, 0x74, 0x6f, 0x70, - 0x69, 0x63, 0x88, 0x01, 0x01, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, - 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1a, 0x0a, - 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x42, 0x05, 0x0a, 0x03, 0x5f, 0x69, 0x64, - 0x42, 0x07, 0x0a, 0x05, 0x5f, 0x6e, 0x61, 0x6d, 0x65, 0x42, 0x08, 0x0a, 0x06, 0x5f, 0x74, 0x6f, - 0x70, 0x69, 0x63, 0x22, 0x76, 0x0a, 0x16, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, - 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x34, 0x0a, - 0x0b, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x18, 0x01, 0x20, 0x03, - 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x6f, 0x6c, 0x6c, - 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x0b, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, - 0x6f, 0x6e, 0x73, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, - 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x93, 0x02, 0x0a, 0x17, - 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x19, 0x0a, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, - 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x88, - 0x01, 0x01, 0x12, 0x17, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, - 0x48, 0x02, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x88, 0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x64, - 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x48, 0x03, - 0x52, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x34, - 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, - 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, - 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x00, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, - 0x64, 0x61, 0x74, 0x61, 0x12, 0x27, 0x0a, 0x0e, 0x72, 0x65, 0x73, 0x65, 0x74, 0x5f, 0x6d, 0x65, - 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x06, 0x20, 0x01, 0x28, 0x08, 0x48, 0x00, 0x52, 0x0d, - 0x72, 0x65, 0x73, 0x65, 0x74, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x42, 0x11, 0x0a, - 0x0f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, - 0x42, 0x08, 0x0a, 0x06, 0x5f, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x42, 0x07, 0x0a, 0x05, 0x5f, 0x6e, - 0x61, 0x6d, 0x65, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, - 0x6e, 0x22, 0x6f, 0x0a, 0x0c, 0x4e, 0x6f, 0x74, 0x69, 0x66, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, - 0x6e, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x03, 0x52, 0x02, 0x69, - 0x64, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, - 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, - 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x03, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, - 0x61, 0x74, 0x75, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, - 0x75, 0x73, 0x32, 0xd6, 0x07, 0x0a, 0x05, 0x53, 0x79, 0x73, 0x44, 0x42, 0x12, 0x49, 0x0a, 0x0e, - 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x12, 0x1d, - 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x44, 0x61, - 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, - 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x52, 0x65, 0x73, - 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x47, 0x65, 0x74, 0x44, 0x61, - 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x12, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, - 0x47, 0x65, 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, - 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x44, - 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, - 0x00, 0x12, 0x45, 0x0a, 0x0c, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x54, 0x65, 0x6e, 0x61, 0x6e, - 0x74, 0x12, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, - 0x65, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, - 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x42, 0x0a, 0x09, 0x47, 0x65, 0x74, 0x54, - 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x18, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, - 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x19, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, - 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x47, 0x0a, 0x0d, - 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x1c, 0x2e, - 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, - 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x63, 0x68, - 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x47, 0x0a, 0x0d, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, - 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x1c, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, - 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, - 0x72, 0x6f, 0x6d, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, - 0x0a, 0x0b, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x1a, 0x2e, - 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, - 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, - 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x47, 0x0a, 0x0d, 0x55, 0x70, 0x64, 0x61, - 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x1c, 0x2e, 0x63, 0x68, 0x72, 0x6f, - 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, - 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, - 0x00, 0x12, 0x57, 0x0a, 0x10, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, - 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x1f, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, - 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, - 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, - 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4d, 0x0a, 0x10, 0x44, 0x65, - 0x6c, 0x65, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x1f, - 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x43, 0x6f, - 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x52, - 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x51, 0x0a, 0x0e, 0x47, 0x65, 0x74, - 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x12, 0x1d, 0x2e, 0x63, 0x68, - 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, - 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x63, 0x68, 0x72, - 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, - 0x6e, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4d, 0x0a, 0x10, - 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, - 0x12, 0x1f, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, - 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x3e, 0x0a, 0x0a, 0x52, - 0x65, 0x73, 0x65, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x12, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, - 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, - 0x79, 0x1a, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x43, 0x5a, 0x41, 0x67, - 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, - 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2d, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, 0x61, - 0x74, 0x6f, 0x72, 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x70, 0x62, - 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, -} - -var ( - file_chromadb_proto_coordinator_proto_rawDescOnce sync.Once - file_chromadb_proto_coordinator_proto_rawDescData = file_chromadb_proto_coordinator_proto_rawDesc -) - -func file_chromadb_proto_coordinator_proto_rawDescGZIP() []byte { - file_chromadb_proto_coordinator_proto_rawDescOnce.Do(func() { - file_chromadb_proto_coordinator_proto_rawDescData = protoimpl.X.CompressGZIP(file_chromadb_proto_coordinator_proto_rawDescData) - }) - return file_chromadb_proto_coordinator_proto_rawDescData -} - -var file_chromadb_proto_coordinator_proto_msgTypes = make([]protoimpl.MessageInfo, 18) -var file_chromadb_proto_coordinator_proto_goTypes = []interface{}{ - (*CreateDatabaseRequest)(nil), // 0: chroma.CreateDatabaseRequest - (*GetDatabaseRequest)(nil), // 1: chroma.GetDatabaseRequest - (*GetDatabaseResponse)(nil), // 2: chroma.GetDatabaseResponse - (*CreateTenantRequest)(nil), // 3: chroma.CreateTenantRequest - (*GetTenantRequest)(nil), // 4: chroma.GetTenantRequest - (*GetTenantResponse)(nil), // 5: chroma.GetTenantResponse - (*CreateSegmentRequest)(nil), // 6: chroma.CreateSegmentRequest - (*DeleteSegmentRequest)(nil), // 7: chroma.DeleteSegmentRequest - (*GetSegmentsRequest)(nil), // 8: chroma.GetSegmentsRequest - (*GetSegmentsResponse)(nil), // 9: chroma.GetSegmentsResponse - (*UpdateSegmentRequest)(nil), // 10: chroma.UpdateSegmentRequest - (*CreateCollectionRequest)(nil), // 11: chroma.CreateCollectionRequest - (*CreateCollectionResponse)(nil), // 12: chroma.CreateCollectionResponse - (*DeleteCollectionRequest)(nil), // 13: chroma.DeleteCollectionRequest - (*GetCollectionsRequest)(nil), // 14: chroma.GetCollectionsRequest - (*GetCollectionsResponse)(nil), // 15: chroma.GetCollectionsResponse - (*UpdateCollectionRequest)(nil), // 16: chroma.UpdateCollectionRequest - (*Notification)(nil), // 17: chroma.Notification - (*Database)(nil), // 18: chroma.Database - (*Status)(nil), // 19: chroma.Status - (*Tenant)(nil), // 20: chroma.Tenant - (*Segment)(nil), // 21: chroma.Segment - (SegmentScope)(0), // 22: chroma.SegmentScope - (*UpdateMetadata)(nil), // 23: chroma.UpdateMetadata - (*Collection)(nil), // 24: chroma.Collection - (*emptypb.Empty)(nil), // 25: google.protobuf.Empty - (*ChromaResponse)(nil), // 26: chroma.ChromaResponse -} -var file_chromadb_proto_coordinator_proto_depIdxs = []int32{ - 18, // 0: chroma.GetDatabaseResponse.database:type_name -> chroma.Database - 19, // 1: chroma.GetDatabaseResponse.status:type_name -> chroma.Status - 20, // 2: chroma.GetTenantResponse.tenant:type_name -> chroma.Tenant - 19, // 3: chroma.GetTenantResponse.status:type_name -> chroma.Status - 21, // 4: chroma.CreateSegmentRequest.segment:type_name -> chroma.Segment - 22, // 5: chroma.GetSegmentsRequest.scope:type_name -> chroma.SegmentScope - 21, // 6: chroma.GetSegmentsResponse.segments:type_name -> chroma.Segment - 19, // 7: chroma.GetSegmentsResponse.status:type_name -> chroma.Status - 23, // 8: chroma.UpdateSegmentRequest.metadata:type_name -> chroma.UpdateMetadata - 23, // 9: chroma.CreateCollectionRequest.metadata:type_name -> chroma.UpdateMetadata - 24, // 10: chroma.CreateCollectionResponse.collection:type_name -> chroma.Collection - 19, // 11: chroma.CreateCollectionResponse.status:type_name -> chroma.Status - 24, // 12: chroma.GetCollectionsResponse.collections:type_name -> chroma.Collection - 19, // 13: chroma.GetCollectionsResponse.status:type_name -> chroma.Status - 23, // 14: chroma.UpdateCollectionRequest.metadata:type_name -> chroma.UpdateMetadata - 0, // 15: chroma.SysDB.CreateDatabase:input_type -> chroma.CreateDatabaseRequest - 1, // 16: chroma.SysDB.GetDatabase:input_type -> chroma.GetDatabaseRequest - 3, // 17: chroma.SysDB.CreateTenant:input_type -> chroma.CreateTenantRequest - 4, // 18: chroma.SysDB.GetTenant:input_type -> chroma.GetTenantRequest - 6, // 19: chroma.SysDB.CreateSegment:input_type -> chroma.CreateSegmentRequest - 7, // 20: chroma.SysDB.DeleteSegment:input_type -> chroma.DeleteSegmentRequest - 8, // 21: chroma.SysDB.GetSegments:input_type -> chroma.GetSegmentsRequest - 10, // 22: chroma.SysDB.UpdateSegment:input_type -> chroma.UpdateSegmentRequest - 11, // 23: chroma.SysDB.CreateCollection:input_type -> chroma.CreateCollectionRequest - 13, // 24: chroma.SysDB.DeleteCollection:input_type -> chroma.DeleteCollectionRequest - 14, // 25: chroma.SysDB.GetCollections:input_type -> chroma.GetCollectionsRequest - 16, // 26: chroma.SysDB.UpdateCollection:input_type -> chroma.UpdateCollectionRequest - 25, // 27: chroma.SysDB.ResetState:input_type -> google.protobuf.Empty - 26, // 28: chroma.SysDB.CreateDatabase:output_type -> chroma.ChromaResponse - 2, // 29: chroma.SysDB.GetDatabase:output_type -> chroma.GetDatabaseResponse - 26, // 30: chroma.SysDB.CreateTenant:output_type -> chroma.ChromaResponse - 5, // 31: chroma.SysDB.GetTenant:output_type -> chroma.GetTenantResponse - 26, // 32: chroma.SysDB.CreateSegment:output_type -> chroma.ChromaResponse - 26, // 33: chroma.SysDB.DeleteSegment:output_type -> chroma.ChromaResponse - 9, // 34: chroma.SysDB.GetSegments:output_type -> chroma.GetSegmentsResponse - 26, // 35: chroma.SysDB.UpdateSegment:output_type -> chroma.ChromaResponse - 12, // 36: chroma.SysDB.CreateCollection:output_type -> chroma.CreateCollectionResponse - 26, // 37: chroma.SysDB.DeleteCollection:output_type -> chroma.ChromaResponse - 15, // 38: chroma.SysDB.GetCollections:output_type -> chroma.GetCollectionsResponse - 26, // 39: chroma.SysDB.UpdateCollection:output_type -> chroma.ChromaResponse - 26, // 40: chroma.SysDB.ResetState:output_type -> chroma.ChromaResponse - 28, // [28:41] is the sub-list for method output_type - 15, // [15:28] is the sub-list for method input_type - 15, // [15:15] is the sub-list for extension type_name - 15, // [15:15] is the sub-list for extension extendee - 0, // [0:15] is the sub-list for field type_name -} - -func init() { file_chromadb_proto_coordinator_proto_init() } -func file_chromadb_proto_coordinator_proto_init() { - if File_chromadb_proto_coordinator_proto != nil { - return - } - file_chromadb_proto_chroma_proto_init() - if !protoimpl.UnsafeEnabled { - file_chromadb_proto_coordinator_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*CreateDatabaseRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetDatabaseRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetDatabaseResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*CreateTenantRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetTenantRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetTenantResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*CreateSegmentRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*DeleteSegmentRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[8].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetSegmentsRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[9].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetSegmentsResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[10].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*UpdateSegmentRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[11].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*CreateCollectionRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[12].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*CreateCollectionResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[13].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*DeleteCollectionRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[14].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetCollectionsRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[15].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*GetCollectionsResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[16].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*UpdateCollectionRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_chromadb_proto_coordinator_proto_msgTypes[17].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*Notification); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - } - file_chromadb_proto_coordinator_proto_msgTypes[8].OneofWrappers = []interface{}{} - file_chromadb_proto_coordinator_proto_msgTypes[10].OneofWrappers = []interface{}{ - (*UpdateSegmentRequest_Topic)(nil), - (*UpdateSegmentRequest_ResetTopic)(nil), - (*UpdateSegmentRequest_Collection)(nil), - (*UpdateSegmentRequest_ResetCollection)(nil), - (*UpdateSegmentRequest_Metadata)(nil), - (*UpdateSegmentRequest_ResetMetadata)(nil), - } - file_chromadb_proto_coordinator_proto_msgTypes[11].OneofWrappers = []interface{}{} - file_chromadb_proto_coordinator_proto_msgTypes[14].OneofWrappers = []interface{}{} - file_chromadb_proto_coordinator_proto_msgTypes[16].OneofWrappers = []interface{}{ - (*UpdateCollectionRequest_Metadata)(nil), - (*UpdateCollectionRequest_ResetMetadata)(nil), - } - type x struct{} - out := protoimpl.TypeBuilder{ - File: protoimpl.DescBuilder{ - GoPackagePath: reflect.TypeOf(x{}).PkgPath(), - RawDescriptor: file_chromadb_proto_coordinator_proto_rawDesc, - NumEnums: 0, - NumMessages: 18, - NumExtensions: 0, - NumServices: 1, - }, - GoTypes: file_chromadb_proto_coordinator_proto_goTypes, - DependencyIndexes: file_chromadb_proto_coordinator_proto_depIdxs, - MessageInfos: file_chromadb_proto_coordinator_proto_msgTypes, - }.Build() - File_chromadb_proto_coordinator_proto = out.File - file_chromadb_proto_coordinator_proto_rawDesc = nil - file_chromadb_proto_coordinator_proto_goTypes = nil - file_chromadb_proto_coordinator_proto_depIdxs = nil -} diff --git a/go/coordinator/migrations/20231129183041.sql b/go/coordinator/migrations/20231129183041.sql deleted file mode 100644 index 2a31ebb4877..00000000000 --- a/go/coordinator/migrations/20231129183041.sql +++ /dev/null @@ -1,8 +0,0 @@ --- Create "notifications" table -CREATE TABLE "public"."notifications" ( - "id" bigserial NOT NULL, - "collection_id" text NULL, - "type" text NULL, - "status" text NULL, - PRIMARY KEY ("id") -); diff --git a/go/coordinator/migrations/atlas.sum b/go/coordinator/migrations/atlas.sum deleted file mode 100644 index d4ee513fa90..00000000000 --- a/go/coordinator/migrations/atlas.sum +++ /dev/null @@ -1,3 +0,0 @@ -h1:j28ectYxexGfQz/LClD7yYVUHAfIcPHlboAJ1Qw0G7I= -20231116210409.sql h1:vwZRvrXrUMOuDykEaheyEzsnNCpmH73x0QEefzUtf8o= -20231129183041.sql h1:FglI5Hjf7kqvjCsSYWkK2IGS2aThQBaVhpg9WekhNEA= diff --git a/go/coordinator/go.mod b/go/go.mod similarity index 70% rename from go/coordinator/go.mod rename to go/go.mod index 93b04935f57..cd45083f25e 100644 --- a/go/coordinator/go.mod +++ b/go/go.mod @@ -1,21 +1,26 @@ -module github.com/chroma/chroma-coordinator +module github.com/chroma-core/chroma/go -go 1.20 +go 1.21 require ( - ariga.io/atlas-provider-gorm v0.1.1 + ariga.io/atlas-provider-gorm v0.3.1 github.com/apache/pulsar-client-go v0.9.1-0.20231030094548-620ecf4addfb - github.com/google/uuid v1.3.1 + github.com/google/uuid v1.6.0 github.com/pingcap/log v1.1.0 github.com/rs/zerolog v1.31.0 github.com/spf13/cobra v1.7.0 github.com/stretchr/testify v1.8.4 + go.opentelemetry.io/otel v1.24.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0 + go.opentelemetry.io/otel/sdk v1.24.0 + go.opentelemetry.io/otel/trace v1.24.0 go.uber.org/automaxprocs v1.5.3 go.uber.org/zap v1.26.0 - google.golang.org/grpc v1.58.3 - google.golang.org/protobuf v1.31.0 + google.golang.org/grpc v1.62.1 + google.golang.org/protobuf v1.33.0 gorm.io/driver/sqlite v1.5.4 - gorm.io/gorm v1.25.5 + gorm.io/gorm v1.25.7 k8s.io/apimachinery v0.28.3 k8s.io/client-go v0.28.3 pgregory.net/rapid v1.1.0 @@ -26,19 +31,26 @@ require ( github.com/99designs/keyring v1.2.1 // indirect github.com/AthenZ/athenz v1.10.39 // indirect github.com/DataDog/zstd v1.5.0 // indirect + github.com/alecthomas/kong v0.7.1 // indirect github.com/ardielle/ardielle-go v1.5.2 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/bits-and-blooms/bitset v1.4.0 // indirect + github.com/cenkalti/backoff/v4 v4.2.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/danieljoos/wincred v1.1.2 // indirect github.com/dvsekhvalnov/jose2go v1.5.0 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 // indirect github.com/golang-jwt/jwt v3.2.1+incompatible // indirect + github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect + github.com/golang-sql/sqlexp v0.1.0 // indirect github.com/golang/snappy v0.0.1 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0 // indirect github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c // indirect github.com/klauspost/compress v1.14.4 // indirect github.com/linkedin/goavro/v2 v2.9.8 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect + github.com/microsoft/go-mssqldb v1.6.0 // indirect github.com/mtibben/percent v0.2.1 // indirect github.com/pierrec/lz4 v2.0.5+incompatible // indirect github.com/prometheus/client_golang v1.11.1 // indirect @@ -46,30 +58,35 @@ require ( github.com/prometheus/common v0.26.0 // indirect github.com/prometheus/procfs v0.6.0 // indirect github.com/sirupsen/logrus v1.8.1 // indirect + go.opentelemetry.io/otel/metric v1.24.0 // indirect + go.opentelemetry.io/proto/otlp v1.1.0 // indirect go.uber.org/atomic v1.9.0 // indirect - golang.org/x/mod v0.11.0 // indirect + golang.org/x/mod v0.15.0 // indirect + golang.org/x/tools v0.17.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240123012728-ef4313101c80 // indirect gorm.io/driver/mysql v1.5.2 // indirect + gorm.io/driver/sqlserver v1.5.2 // indirect ) require ( - ariga.io/atlas-go-sdk v0.1.1-0.20231001054405-7edfcfc14f1c // indirect + ariga.io/atlas-go-sdk v0.2.3 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/emicklei/go-restful/v3 v3.9.0 // indirect github.com/evanphx/json-patch v4.12.0+incompatible // indirect - github.com/go-logr/logr v1.2.4 // indirect + github.com/go-logr/logr v1.4.1 // indirect github.com/go-openapi/jsonpointer v0.19.6 // indirect github.com/go-openapi/jsonreference v0.20.2 // indirect github.com/go-openapi/swag v0.22.3 // indirect github.com/go-sql-driver/mysql v1.7.1 // indirect github.com/gogo/protobuf v1.3.2 // indirect - github.com/golang/protobuf v1.5.3 // indirect + github.com/golang/protobuf v1.5.4 // indirect github.com/google/gnostic-models v0.6.8 // indirect - github.com/google/go-cmp v0.5.9 // indirect + github.com/google/go-cmp v0.6.0 // indirect github.com/google/gofuzz v1.2.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect - github.com/jackc/pgx/v5 v5.3.1 // indirect + github.com/jackc/pgx/v5 v5.3.1 github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/josharian/intern v1.0.0 // indirect @@ -88,15 +105,15 @@ require ( github.com/spf13/pflag v1.0.5 // indirect github.com/stretchr/objx v0.5.1 // indirect go.uber.org/multierr v1.11.0 // indirect - golang.org/x/crypto v0.15.0 // indirect - golang.org/x/net v0.18.0 // indirect - golang.org/x/oauth2 v0.10.0 // indirect - golang.org/x/sys v0.14.0 // indirect - golang.org/x/term v0.14.0 // indirect + golang.org/x/crypto v0.21.0 // indirect + golang.org/x/net v0.22.0 // indirect + golang.org/x/oauth2 v0.16.0 // indirect + golang.org/x/sys v0.18.0 // indirect + golang.org/x/term v0.18.0 // indirect golang.org/x/text v0.14.0 // indirect golang.org/x/time v0.3.0 // indirect - google.golang.org/appengine v1.6.7 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20231016165738-49dd2c1f3d0b // indirect + google.golang.org/appengine v1.6.8 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240325203815-454cdb8f5daa gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect gopkg.in/yaml.v2 v2.4.0 // indirect diff --git a/go/coordinator/go.sum b/go/go.sum similarity index 67% rename from go/coordinator/go.sum rename to go/go.sum index 15390626451..9bc3edc2f27 100644 --- a/go/coordinator/go.sum +++ b/go/go.sum @@ -1,7 +1,7 @@ -ariga.io/atlas-go-sdk v0.1.1-0.20231001054405-7edfcfc14f1c h1:jvi4KB/7DmYYT+Wy2TFImccaBU0+dw7V8Un67NDGuio= -ariga.io/atlas-go-sdk v0.1.1-0.20231001054405-7edfcfc14f1c/go.mod h1:MLvZ9QwZx1KhI6+8XguxHPUPm0/PTTUr46S5GQAe9WI= -ariga.io/atlas-provider-gorm v0.1.1 h1:Y0VsZCQkXJRYIJxenn2BM6sW2u9SkTca5mLvJumqrgE= -ariga.io/atlas-provider-gorm v0.1.1/go.mod h1:jb8uYcN+ul8Nf7OVzi5Vd2y+SQXrI4dHYBEUCiCi/6Q= +ariga.io/atlas-go-sdk v0.2.3 h1:DpKruiJ9ElJcNhYxnQM9ddzupHXEYFH0Jx6ZcZ7lKYQ= +ariga.io/atlas-go-sdk v0.2.3/go.mod h1:owkEEXw6jqne5KPVDfKsYB7cwMiMk3jtOiAAeKxS/yU= +ariga.io/atlas-provider-gorm v0.3.1 h1:+RrnoBwlqMj+B1x/Cf1BfwtZzq6v5vKzHdl2A6nZuBU= +ariga.io/atlas-provider-gorm v0.3.1/go.mod h1:NOXGkyHfWFm8vQO7T+je5Zj5DdLZhkzReXGfxnnK4VM= cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 h1:/vQbFIOMbk2FiG/kXiLl8BRyzTWDw7gX/Hz7Dd5eDMs= github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4/go.mod h1:hN7oaIRCjzsZ2dE+yG5k+rsdt3qcwykqK6HVGcKwsw4= @@ -9,9 +9,29 @@ github.com/99designs/keyring v1.2.1 h1:tYLp1ULvO7i3fI5vE21ReQuj99QFSs7lGm0xWyJo8 github.com/99designs/keyring v1.2.1/go.mod h1:fc+wB5KTk9wQ9sDx0kFXB3A0MaeGHM9AwRStKOQ5vOA= github.com/AthenZ/athenz v1.10.39 h1:mtwHTF/v62ewY2Z5KWhuZgVXftBej1/Tn80zx4DcawY= github.com/AthenZ/athenz v1.10.39/go.mod h1:3Tg8HLsiQZp81BJY58JBeU2BR6B/H4/0MQGfCwhHNEA= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.4.0/go.mod h1:ON4tFdPTwRcgWEaVDrN3584Ef+b7GgSJaXxe5fW9t4M= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.6.0/go.mod h1:bjGvMhVMb+EEm3VRNQawDMUyMMjo+S5ewNjflkep/0Q= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.6.1/go.mod h1:bjGvMhVMb+EEm3VRNQawDMUyMMjo+S5ewNjflkep/0Q= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.7.1 h1:/iHxaJhsFr0+xVFfbMr5vxz848jyiWuIEDhYq3y5odY= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.7.1/go.mod h1:bjGvMhVMb+EEm3VRNQawDMUyMMjo+S5ewNjflkep/0Q= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.3.0 h1:vcYCAze6p19qBW7MhZybIsqD8sMV8js0NyQM8JDnVtg= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.3.0/go.mod h1:OQeznEEkTZ9OrhHJoDD8ZDq51FHgXjqtP9z6bEwBq9U= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.2/go.mod h1:eWRD7oawr1Mu1sLCawqVc0CUiF43ia3qQMxLscsKQ9w= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.2.0/go.mod h1:eWRD7oawr1Mu1sLCawqVc0CUiF43ia3qQMxLscsKQ9w= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.3.0 h1:sXr+ck84g/ZlZUOZiNELInmMgOsuGwdjjVkEIde0OtY= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.3.0/go.mod h1:okt5dMMTOFjX/aovMlrjvvXoPMBVSPzk9185BT0+eZM= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys v1.0.0 h1:yfJe15aSwEQ6Oo6J+gdfdulPNoZ3TEhmbhLIoxZcA+U= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys v1.0.0/go.mod h1:Q28U+75mpCaSCDowNEmhIo/rmgdkqmkmzI7N6TGR4UY= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal v0.8.0 h1:T028gtTPiYt/RMUfs8nVsAL7FDQrfLlrm/NnRG/zcC4= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal v0.8.0/go.mod h1:cw4zVQgBby0Z5f2v0itn6se2dDP17nTjbZFXW5uPyHA= +github.com/AzureAD/microsoft-authentication-library-for-go v1.0.0/go.mod h1:kgDmCTgBzIEPFElEF+FK0SdjAor06dRq2Go927dnQ6o= +github.com/AzureAD/microsoft-authentication-library-for-go v1.1.0 h1:HCc0+LpPfpCKs6LGGLAhwBARt9632unrVcI6i8s/8os= +github.com/AzureAD/microsoft-authentication-library-for-go v1.1.0/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/DataDog/zstd v1.5.0 h1:+K/VEwIAaPcHiMtQvpLD4lqW7f0Gk3xdYZmI1hD+CXo= github.com/DataDog/zstd v1.5.0/go.mod h1:g4AWEaM3yOg3HYfnJ3YIawPnVdXJh9QME85blwSAmyw= +github.com/alecthomas/kong v0.7.1 h1:azoTh0IOfwlAX3qN9sHWTxACE2oV8Bg2gAwBsMwDQY4= +github.com/alecthomas/kong v0.7.1/go.mod h1:n1iCIO2xS46oE8ZfYCNDqdR0b0wZNrXAIAqro/2132U= github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc= github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc= github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0= @@ -30,6 +50,8 @@ github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/bits-and-blooms/bitset v1.4.0 h1:+YZ8ePm+He2pU3dZlIZiOeAKfrBkXi1lSrXJ/Xzgbu8= github.com/bits-and-blooms/bitset v1.4.0/go.mod h1:gIdJ4wp64HaoK2YrL1Q5/N7Y16edYb8uY+O0FJTyyDA= +github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM= +github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= @@ -43,6 +65,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dimfeld/httptreemux v5.0.1+incompatible h1:Qj3gVcDNoOthBAqftuD596rm4wg/adLLz5xh5CmpiCA= github.com/dimfeld/httptreemux v5.0.1+incompatible/go.mod h1:rbUlSV+CCpv/SuqUTP/8Bk2O3LyUV436/yaRGkhP6Z0= +github.com/dnaeon/go-vcr v1.1.0/go.mod h1:M7tiix8f0r6mKKJ3Yq/kqU1OYf3MnfmBWVbPx/yU9ko= +github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ= github.com/dvsekhvalnov/jose2go v1.5.0 h1:3j8ya4Z4kMCwT5nXIKFSV84YS+HdqSSO0VsTQxaLAeM= github.com/dvsekhvalnov/jose2go v1.5.0/go.mod h1:QsHjhyTlD/lAVqn/NSbVZmSCGeDehTB/mPZadG+mhXU= github.com/emicklei/go-restful/v3 v3.9.0 h1:XwGDlfxEnQZzuopoqxwSEllNcCOM9DhhFyhFIIGKwxE= @@ -50,6 +74,7 @@ github.com/emicklei/go-restful/v3 v3.9.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry github.com/evanphx/json-patch v4.12.0+incompatible h1:4onqiflcdA9EOZ4RxV643DvftH5pOlLGNtQ5lPWQu84= github.com/evanphx/json-patch v4.12.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk= github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as= github.com/go-kit/kit v0.9.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as= github.com/go-kit/log v0.1.0/go.mod h1:zbhenjAZHb184qTLMA9ZjW7ThYL0H2mk7Q6pNt4vbaY= @@ -57,8 +82,11 @@ github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9 github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk= github.com/go-logfmt/logfmt v0.5.0/go.mod h1:wCYkCAKZfumFQihp8CzCvQ3paCTfi41vtzG1KdI/P7A= github.com/go-logr/logr v1.2.0/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= -github.com/go-logr/logr v1.2.4 h1:g01GSCwiDw2xSZfjJ2/T9M+S6pFdcNtFYsp+Y43HYDQ= -github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.1 h1:pKouT5E8xu9zeFC39JXRDukb6JFQPXM5p5I91188VAQ= +github.com/go-logr/logr v1.4.1/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/go-openapi/jsonpointer v0.19.6 h1:eCs3fxoIi3Wh6vtgmLTOjdhSpiqphQ+DaPn38N2ZdrE= github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= github.com/go-openapi/jsonreference v0.20.2 h1:3sVjiK66+uXK/6oQ8xgcRKcFgQ5KXa2KvnJRumpMGbE= @@ -71,6 +99,7 @@ github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrt github.com/go-sql-driver/mysql v1.7.1/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI= +github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 h1:ZpnhV/YsD2/4cESfV5+Hoeu/iUR3ruzNvZ+yQfO03a0= github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2/go.mod h1:bBOAhwG1umN6/6ZUMtDFBMQR8jRg9O75tm9K00oMsK4= github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= @@ -79,6 +108,14 @@ github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= github.com/golang-jwt/jwt v3.2.1+incompatible h1:73Z+4BJcrTC+KczS6WvTPvRGOp1WmfEP4Q1lOd9Z/+c= github.com/golang-jwt/jwt v3.2.1+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I= +github.com/golang-jwt/jwt/v4 v4.4.3/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= +github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= +github.com/golang-jwt/jwt/v5 v5.0.0 h1:1n1XNM9hk7O9mnQoNBGolZvzebBQ7p93ULHRc28XJUE= +github.com/golang-jwt/jwt/v5 v5.0.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= +github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA= +github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0= +github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei6A= +github.com/golang-sql/sqlexp v0.1.0/go.mod h1:J4ad9Vo8ZCWQ2GMrC4UCQy1JpCbwU9m3EOqtpKwwwHI= github.com/golang/mock v1.4.3/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= @@ -91,8 +128,11 @@ github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvq github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4= github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/gnostic-models v0.6.8 h1:yo/ABAfM5IMRsS1VnXjTBvUb61tFIHozhlYvRgGre9I= @@ -103,17 +143,27 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.1/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= -github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0= github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1 h1:K6RDEckDVWvDI9JAJYCmNdQXq6neHJOYx3V6jnqNEec= -github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4= -github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4= +github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/mux v1.7.4/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/gorilla/securecookie v1.1.1/go.mod h1:ra0sb63/xPlUeL+yeDciTfxMRAA+MP+HVt/4epWDjd4= +github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0 h1:Wqo399gCIufwto+VfwCSvsnfGpF/w5E9CNxSwbpD6No= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0/go.mod h1:qmOFXW2epJhM0qSnUUYpldc7gVz2KMQwJ/QYCDIa7XU= github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c h1:6rhixN/i8ZofjG1Y75iExal34USq5p+wiN1tpie8IrU= github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c/go.mod h1:NMPJylDgVpX0MLRlPy15sqSwOFv/U1GZ2m21JhFfek0= +github.com/hashicorp/go-uuid v1.0.2/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= +github.com/hashicorp/go-uuid v1.0.3/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= @@ -124,6 +174,12 @@ github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU= github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8= github.com/jawher/mow.cli v1.0.4/go.mod h1:5hQj2V8g+qYmLUVWqu4Wuja1pI57M83EChYLVZ0sMKk= github.com/jawher/mow.cli v1.2.0/go.mod h1:y+pcA3jBAdo/GIZx/0rFjw/K2bVEODP9rfZOfaiq8Ko= +github.com/jcmturner/aescts/v2 v2.0.0/go.mod h1:AiaICIRyfYg35RUkr8yESTqvSy7csK90qZ5xfvvsoNs= +github.com/jcmturner/dnsutils/v2 v2.0.0/go.mod h1:b0TnjGOvI/n42bZa+hmXL+kFJZsFT7G4t3HTlQ184QM= +github.com/jcmturner/gofork v1.7.6/go.mod h1:1622LH6i/EZqLloHfE7IeZ0uEJwMSUyQ/nDd82IeqRo= +github.com/jcmturner/goidentity/v6 v6.0.1/go.mod h1:X1YW3bgtvwAXju7V3LCIMpY0Gbxyjn/mY9zx4tFonSg= +github.com/jcmturner/gokrb5/v8 v8.4.4/go.mod h1:1btQEpgT6k+unzCwX1KdWMEwPPkkgBtP+F6aCACiMrs= +github.com/jcmturner/rpc/v2 v2.0.3/go.mod h1:VUJYCIDm3PVOEHw8sgt091/20OJjskO/YJki3ELg/Hc= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= @@ -150,10 +206,16 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/lib/pq v1.10.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= +github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= +github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/linkedin/goavro/v2 v2.9.8 h1:jN50elxBsGBDGVDEKqUlDuU1cFwJ11K/yrJCBMe/7Wg= github.com/linkedin/goavro/v2 v2.9.8/go.mod h1:UgQUb2N/pmueQYH9bfqFioWxzYCZXSfF8Jw03O5sjqA= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= @@ -167,6 +229,8 @@ github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6 github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= +github.com/microsoft/go-mssqldb v1.6.0 h1:mM3gYdVwEPFrlg/Dvr2DNVEgYFG7L42l+dGc67NNNpc= +github.com/microsoft/go-mssqldb v1.6.0/go.mod h1:00mDtPbeQCRGC1HwOOR5K/gr30P1NcEG0vx6Kbv2aJU= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -174,6 +238,8 @@ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lN github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0= github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/modocache/gover v0.0.0-20171022184752-b58185e213c5/go.mod h1:caMODM3PzxT8aQXRPkAt8xlV/e7d7w8GM5g0fa5F0D8= +github.com/montanaflynn/stats v0.7.0/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow= github.com/mtibben/percent v0.2.1 h1:5gssi8Nqo8QU/r2pynCm+hBQHpkB/uNK7BJCFogWdzs= github.com/mtibben/percent v0.2.1/go.mod h1:KG9uO+SZkUp+VkRHsCdYQV3XSZrrSpR3O9ibNBTZrns= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= @@ -182,15 +248,21 @@ github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRW github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno= github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= +github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU= github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= +github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU= github.com/onsi/ginkgo/v2 v2.9.4 h1:xR7vG4IXt5RWx6FfIjyAtsoMAtnc3C/rFXBBd2AjZwE= +github.com/onsi/ginkgo/v2 v2.9.4/go.mod h1:gCQYp2Q+kSoIj7ykSVb9nskRSsR6PUj4AiLywzIhbKM= github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE= +github.com/onsi/gomega v1.27.6/go.mod h1:PIQNjfQwkP3aQAH7lf7j87O/5FiNr+ZR8+ipb+qQlhg= github.com/pierrec/lz4 v2.0.5+incompatible h1:2xWsjqPFWcplujydGg4WmhC/6fZqK42wMM8aXeqhl0I= github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY= +github.com/pingcap/errors v0.11.0 h1:DCJQB8jrHbQ1VVlMFIrbj2ApScNNotVmkSNplu2yUt4= github.com/pingcap/errors v0.11.0/go.mod h1:Oi8TUi2kEtXXLMJk9l1cGmz20kV3TaQ0usTwv5KuLY8= -github.com/pingcap/errors v0.11.5-0.20210425183316-da1aaba5fb63 h1:+FZIDR/D97YOPik4N4lPDaUcLDF/EQPogxtlHB2ZZRM= github.com/pingcap/log v1.1.0 h1:ELiPxACz7vdo1qAvvaWJg1NrYFoY6gqAh/+Uo6aXdD8= github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4= +github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 h1:KoWmjvw+nsYOo29YJK9vDA65RGE3NrOnUtO7a+RF9HU= +github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8/go.mod h1:HKlIX3XHQyzLZPlr7++PzdhaXEj94dEiJgZDTsxEqUI= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= @@ -198,6 +270,7 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g= +github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U= github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw= github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo= github.com/prometheus/client_golang v1.7.1/go.mod h1:PY5Wy2awLA44sXw4AOSfFBetzPP4j5+D6mVACh+pe2M= @@ -253,13 +326,29 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= +github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= +go.opentelemetry.io/otel v1.24.0 h1:0LAOdjNmQeSTzGBzduGe/rU4tZhMwL5rWgtp9Ku5Jfo= +go.opentelemetry.io/otel v1.24.0/go.mod h1:W7b9Ozg4nkF5tWI5zsXkaKKDjdVjpD4oAt9Qi/MArHo= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0 h1:t6wl9SPayj+c7lEIFgm4ooDBZVb01IhLB4InpomhRw8= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0/go.mod h1:iSDOcsnSA5INXzZtwaBPrKp/lWu/V14Dd+llD0oI2EA= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0 h1:Mw5xcxMwlqoJd97vwPxA8isEaIoxsta9/Q51+TTJLGE= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0/go.mod h1:CQNu9bj7o7mC6U7+CA/schKEYakYXWr79ucDHTMGhCM= +go.opentelemetry.io/otel/metric v1.24.0 h1:6EhoGWWK28x1fbpA4tYTOWBkPefTDQnb8WSGXlc88kI= +go.opentelemetry.io/otel/metric v1.24.0/go.mod h1:VYhLe1rFfxuTXLgj4CBiyz+9WYBA8pNGJgDcSFRKBco= +go.opentelemetry.io/otel/sdk v1.24.0 h1:YMPPDNymmQN3ZgczicBY3B6sf9n62Dlj9pWD3ucgoDw= +go.opentelemetry.io/otel/sdk v1.24.0/go.mod h1:KVrIYw6tEubO9E96HQpcmpTKDVn9gdv35HoYiQWGDFg= +go.opentelemetry.io/otel/trace v1.24.0 h1:CsKnnL4dUAr/0llH9FKuc698G04IrpWV0MQA/Y1YELI= +go.opentelemetry.io/otel/trace v1.24.0/go.mod h1:HPc3Xr/cOApsBI154IU0OI0HJexz+aw5uPdbs3UCjNU= +go.opentelemetry.io/proto/otlp v1.1.0 h1:2Di21piLrCqJ3U3eXGCTPHE9R8Nh+0uglSnOyxikMeI= +go.opentelemetry.io/proto/otlp v1.1.0/go.mod h1:GpBHCBWiqvVLDqmHZsoMM3C5ySeKTC7ej/RNTae6MdY= go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE= go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/automaxprocs v1.5.3 h1:kWazyxZUrS3Gs4qUpbwo5kEIMGe/DAvi5Z4tl2NW4j8= go.uber.org/automaxprocs v1.5.3/go.mod h1:eRbA25aqJrxAbsLO0xy5jVwPt7FQnRgjW+efnwa1WM0= go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A= -go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU= go.uber.org/multierr v1.7.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= @@ -272,37 +361,60 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= -golang.org/x/crypto v0.15.0 h1:frVn1TEaCEaZcn3Tmd7Y2b5KKPaZ+I32Q2OA3kYp5TA= -golang.org/x/crypto v0.15.0/go.mod h1:4ChreQoLWfG3xLDer1WdlH5NdlQ3+mwnQq1YTKY+72g= +golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= +golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58= +golang.org/x/crypto v0.7.0/go.mod h1:pYwdfH91IfpZVANVyUOhSIPZaFoJGxTFbZhFTx+dXZU= +golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0= +golang.org/x/crypto v0.12.0/go.mod h1:NF0Gs7EO5K4qLn+Ylc+fih8BSTeIjAP05siRnAh98yw= +golang.org/x/crypto v0.18.0 h1:PGVlW0xEltQnzFZ55hkuX5+KLyrMYhHld1YHO4AKcdc= +golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1mg= +golang.org/x/crypto v0.21.0 h1:X31++rzVUdKhX5sWmSOFZxx8UW/ldWx55cbf08iNAMA= +golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs= golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= -golang.org/x/mod v0.11.0 h1:bUO06HqtnRcc/7l71XBe4WcqTZ+3AH1J59zWDDwLKgU= -golang.org/x/mod v0.11.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= +golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= +golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= +golang.org/x/mod v0.15.0 h1:SernR4v+D55NyBH2QiEQrlBAnj1ECL6AGrA5+dPaMY8= +golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20181114220301-adae6a3d119a/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190108225652-1e06a53dbb7e/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= -golang.org/x/net v0.0.0-20190603091049-60506f45cf65/go.mod h1:HSz+uSET+XFnRR8LxR5pz3Of3rY3CfYBVs4xY44aLks= golang.org/x/net v0.0.0-20190613194153-d28f0bde5980/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200114155413-6afb5195e5aa/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200202094626-16171245cfb2/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200625001655-4c5254603344/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA= +golang.org/x/net v0.0.0-20201010224723-4f7140c49acb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210726213435-c6fcb2dbf985/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= -golang.org/x/net v0.18.0 h1:mIYleuAkSbHh0tCv7RvjL3F6ZVbLjq4+R7zbOn3Kokg= -golang.org/x/net v0.18.0/go.mod h1:/czyP5RqHAH4odGYxBJ1qz0+CE5WZ+2j1YgoEo8F2jQ= +golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= +golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= +golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= +golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc= +golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= +golang.org/x/net v0.14.0/go.mod h1:PpSgVXXLK0OxS0F31C1/tv6XNguvCrnXIDrFMspZIUI= +golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo= +golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY= +golang.org/x/net v0.22.0 h1:9sGLhx7iRIHEiX0oAJ3MRZMUCElJgy7Br1nO+AMN3Tc= +golang.org/x/net v0.22.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= -golang.org/x/oauth2 v0.10.0 h1:zHCpF2Khkwy4mMB4bv0U37YtJdTGW8jI0glAApi0Kh8= -golang.org/x/oauth2 v0.10.0/go.mod h1:kTpgurOux7LqtuxjuyZa4Gj2gdezIt/jQtGnNFfypQI= +golang.org/x/oauth2 v0.15.0 h1:s8pnnxNVzjWyrvYdFUQq5llS1PX2zhPXmccZv99h7uQ= +golang.org/x/oauth2 v0.15.0/go.mod h1:q48ptWNTY5XWf+JNten23lcvHpLJ0ZSxF5ttTHKVCAM= +golang.org/x/oauth2 v0.16.0 h1:aDkGMBSYxElaoP81NpoUoz2oo2R2wHdZpGToUxfyQrQ= +golang.org/x/oauth2 v0.16.0/go.mod h1:hqZ+0LWXsiVoZpeld6jVt06P3adbS2Uu911W1SsJv2o= golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -318,20 +430,42 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210603081109-ebe580a85c40/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210616045830-e2b7044e8c71/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210819135213-f52c844e1c1c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.14.0 h1:Vz7Qs629MkJkGyHxUlRHizWJRG2j8fbQKjELVSNhy7Q= -golang.org/x/sys v0.14.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y= +golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4= +golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= -golang.org/x/term v0.14.0 h1:LGK9IlZ8T9jvdy6cTdfKUCltatMFOehAQo9SRC46UQ8= -golang.org/x/term v0.14.0/go.mod h1:TySc+nGkYR6qt8km8wUhuFRTVSMIX3XPR58y2lC8vww= +golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= +golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= +golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U= +golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo= +golang.org/x/term v0.11.0/go.mod h1:zC9APTIj3jG3FdV/Ons+XE1riIZXG4aZ4GTHiPZJPIU= +golang.org/x/term v0.16.0 h1:m+B6fahuftsE9qjo0VWp2FW0mB3MTJvR0BaMQrq0pmE= +golang.org/x/term v0.16.0/go.mod h1:yn7UURbUtPyrVJPGPq404EukNFxcm/foM+bV/bfcDsY= +golang.org/x/term v0.18.0 h1:FcHjZXDMxI8mM3nwhX9HlKop4C0YQvCVCdwYl2wOtE8= +golang.org/x/term v0.18.0/go.mod h1:ILwASektA3OnRv7amZ1xhE/KTR+u50pbXfZ03+6Nx58= golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= +golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ= +golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= +golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= +golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= +golang.org/x/text v0.12.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4= @@ -343,18 +477,32 @@ golang.org/x/tools v0.0.0-20191108193012-7d206e10da11/go.mod h1:b+2E5dAYhXwXZwtn golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= -golang.org/x/tools v0.10.0 h1:tvDr/iQoUqNdohiYm0LmmKcBk+q86lb9EprIUFhHHGg= +golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= +golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= +golang.org/x/tools v0.17.0 h1:FvmRgNOcs3kOa+T20R1uhfP9F6HgG2mfxDv1vrx1Htc= +golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= -google.golang.org/appengine v1.6.7 h1:FZR1q0exgwxzPzp/aF+VccGrSfxfPpkBqjIIEq3ru6c= -google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc= -google.golang.org/genproto/googleapis/rpc v0.0.0-20231016165738-49dd2c1f3d0b h1:ZlWIi1wSK56/8hn4QcBp/j9M7Gt3U/3hZw3mC7vDICo= -google.golang.org/genproto/googleapis/rpc v0.0.0-20231016165738-49dd2c1f3d0b/go.mod h1:swOH3j0KzcDDgGUWr+SNpyTen5YrXjS3eyPzFYKc6lc= -google.golang.org/grpc v1.58.3 h1:BjnpXut1btbtgN/6sp+brB2Kbm2LjNXnidYujAVbSoQ= -google.golang.org/grpc v1.58.3/go.mod h1:tgX3ZQDlNJGU96V6yHh1T/JeoBQ2TXdr43YbYSsCJk0= +google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM= +google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds= +google.golang.org/genproto v0.0.0-20231212172506-995d672761c0 h1:YJ5pD9rF8o9Qtta0Cmy9rdBwkSjrTCT6XTiUQVOtIos= +google.golang.org/genproto v0.0.0-20231212172506-995d672761c0/go.mod h1:l/k7rMz0vFTBPy+tFSGvXEd3z+BcoG1k7EHbqm+YBsY= +google.golang.org/genproto v0.0.0-20240123012728-ef4313101c80 h1:KAeGQVN3M9nD0/bQXnr/ClcEMJ968gUXJQ9pwfSynuQ= +google.golang.org/genproto/googleapis/api v0.0.0-20240102182953-50ed04b92917 h1:rcS6EyEaoCO52hQDupoSfrxI3R6C2Tq741is7X8OvnM= +google.golang.org/genproto/googleapis/api v0.0.0-20240102182953-50ed04b92917/go.mod h1:CmlNWB9lSezaYELKS5Ym1r44VrrbPUa7JTvw+6MbpJ0= +google.golang.org/genproto/googleapis/api v0.0.0-20240123012728-ef4313101c80 h1:Lj5rbfG876hIAYFjqiJnPHfhXbv+nzTWfm04Fg/XSVU= +google.golang.org/genproto/googleapis/api v0.0.0-20240123012728-ef4313101c80/go.mod h1:4jWUdICTdgc3Ibxmr8nAJiiLHwQBY0UI0XZcEMaFKaA= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240102182953-50ed04b92917 h1:6G8oQ016D88m1xAKljMlBOOGWDZkes4kMhgGFlf8WcQ= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240102182953-50ed04b92917/go.mod h1:xtjpI3tXFPP051KaWnhvxkiubL/6dJ18vLVf7q2pTOU= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240325203815-454cdb8f5daa h1:RBgMaUMP+6soRkik4VoN8ojR2nex2TqZwjSSogic+eo= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240325203815-454cdb8f5daa/go.mod h1:WtryC6hu0hhx87FDGxWCDptyssuo68sk10vYjF+T9fY= +google.golang.org/grpc v1.61.1 h1:kLAiWrZs7YeDM6MumDe7m3y4aM6wacLzM1Y/wiLP9XY= +google.golang.org/grpc v1.61.1/go.mod h1:VUbo7IFqmF1QtCAstipjG0GIoq49KvMe9+h1jFLBNJs= +google.golang.org/grpc v1.62.1 h1:B4n+nfKzOICUXMgyrNd19h/I9oH0L1pizfk1d4zSgTk= +google.golang.org/grpc v1.62.1/go.mod h1:IWTG0VlJLCh1SkC58F7np9ka9mx/WNkjl4PGJaiq+QE= google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= @@ -363,8 +511,10 @@ google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzi google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= -google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8= -google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +google.golang.org/protobuf v1.32.0 h1:pPC6BG5ex8PDFnkbrGU3EixyhKcQ2aDuBS36lqK/C7I= +google.golang.org/protobuf v1.32.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI= +google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= @@ -379,6 +529,7 @@ gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= gopkg.in/square/go-jose.v2 v2.4.1/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI= gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= @@ -397,9 +548,12 @@ gorm.io/driver/postgres v1.5.2 h1:ytTDxxEv+MplXOfFe3Lzm7SjG09fcdb3Z/c056DTBx0= gorm.io/driver/postgres v1.5.2/go.mod h1:fmpX0m2I1PKuR7mKZiEluwrP3hbs+ps7JIGMUBpCgl8= gorm.io/driver/sqlite v1.5.4 h1:IqXwXi8M/ZlPzH/947tn5uik3aYQslP9BVveoax0nV0= gorm.io/driver/sqlite v1.5.4/go.mod h1:qxAuCol+2r6PannQDpOP1FP6ag3mKi4esLnB/jHed+4= +gorm.io/driver/sqlserver v1.5.2 h1:+o4RQ8w1ohPbADhFqDxeeZnSWjwOcBnxBckjTbcP4wk= +gorm.io/driver/sqlserver v1.5.2/go.mod h1:gaKF0MO0cfTq9Q3/XhkowSw4g6nIwHPGAs4hzKCmvBo= gorm.io/gorm v1.25.2-0.20230530020048-26663ab9bf55/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k= -gorm.io/gorm v1.25.5 h1:zR9lOiiYf09VNh5Q1gphfyia1JpiClIWG9hQaxB/mls= -gorm.io/gorm v1.25.5/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= +gorm.io/gorm v1.25.2-0.20230610234218-206613868439/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k= +gorm.io/gorm v1.25.7 h1:VsD6acwRjz2zFxGO50gPO6AkNs7KKnvfzUjHQhZDz/A= +gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= k8s.io/api v0.28.3 h1:Gj1HtbSdB4P08C8rs9AR94MfSGpRhJgsS+GF9V26xMM= k8s.io/api v0.28.3/go.mod h1:MRCV/jr1dW87/qJnZ57U5Pak65LGmQVkKTzf3AtKFHc= k8s.io/apimachinery v0.28.3 h1:B1wYx8txOaCQG0HmYF6nbpU8dg6HvA06x5tEffvOe7A= diff --git a/go/coordinator/migrations/20231116210409.sql b/go/migrations/20240313233558.sql similarity index 76% rename from go/coordinator/migrations/20231116210409.sql rename to go/migrations/20240313233558.sql index bb9c8d8a00c..0503ce8984e 100644 --- a/go/coordinator/migrations/20231116210409.sql +++ b/go/migrations/20240313233558.sql @@ -1,3 +1,5 @@ +CREATE SCHEMA IF NOT EXISTS "public"; + -- Create "collection_metadata" table CREATE TABLE "public"."collection_metadata" ( "collection_id" text NOT NULL, @@ -10,6 +12,7 @@ CREATE TABLE "public"."collection_metadata" ( "updated_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY ("collection_id", "key") ); + -- Create "collections" table CREATE TABLE "public"."collections" ( "id" text NOT NULL, @@ -21,10 +24,14 @@ CREATE TABLE "public"."collections" ( "is_deleted" boolean NULL DEFAULT false, "created_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, "updated_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + "log_position" bigint NULL DEFAULT 0, + "version" integer NULL DEFAULT 0, PRIMARY KEY ("id") ); --- Create index "collections_name_key" to table: "collections" -CREATE UNIQUE INDEX "collections_name_key" ON "public"."collections" ("name"); + +-- Create index "uni_collections_name" to table: "collections" +CREATE UNIQUE INDEX "uni_collections_name" ON "public"."collections" ("name"); + -- Create "databases" table CREATE TABLE "public"."databases" ( "id" text NOT NULL, @@ -36,8 +43,28 @@ CREATE TABLE "public"."databases" ( "updated_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY ("id") ); + -- Create index "idx_tenantid_name" to table: "databases" CREATE UNIQUE INDEX "idx_tenantid_name" ON "public"."databases" ("name", "tenant_id"); + +-- Create "notifications" table +CREATE TABLE "public"."notifications" ( + "id" bigserial NOT NULL, + "collection_id" text NULL, + "type" text NULL, + "status" text NULL, + PRIMARY KEY ("id") +); + +-- Create "record_logs" table +CREATE TABLE "public"."record_logs" ( + "collection_id" text NOT NULL, + "id" bigint NOT NULL, + "timestamp" bigint NULL, + "record" bytea NULL, + PRIMARY KEY ("collection_id", "id") +); + -- Create "segment_metadata" table CREATE TABLE "public"."segment_metadata" ( "segment_id" text NOT NULL, @@ -50,8 +77,10 @@ CREATE TABLE "public"."segment_metadata" ( "updated_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY ("segment_id", "key") ); + -- Create "segments" table CREATE TABLE "public"."segments" ( + "collection_id" text NOT NULL, "id" text NOT NULL, "type" text NOT NULL, "scope" text NULL, @@ -60,9 +89,10 @@ CREATE TABLE "public"."segments" ( "is_deleted" boolean NULL DEFAULT false, "created_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, "updated_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, - "collection_id" text NULL, - PRIMARY KEY ("id") + "file_paths" text NULL DEFAULT '{}', + PRIMARY KEY ("collection_id", "id") ); + -- Create "tenants" table CREATE TABLE "public"."tenants" ( "id" text NOT NULL, @@ -70,5 +100,6 @@ CREATE TABLE "public"."tenants" ( "is_deleted" boolean NULL DEFAULT false, "created_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, "updated_at" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + "last_compaction_time" bigint NOT NULL, PRIMARY KEY ("id") -); +); \ No newline at end of file diff --git a/go/migrations/20240321194713.sql b/go/migrations/20240321194713.sql new file mode 100644 index 00000000000..c0a778c530d --- /dev/null +++ b/go/migrations/20240321194713.sql @@ -0,0 +1,14 @@ +INSERT INTO + "public"."tenants" (id, last_compaction_time) +VALUES + ('default_tenant', 0); + +-- The default tenant id is 'default_tenant' others are UUIDs +INSERT INTO + databases (id, name, tenant_id) +VALUES + ( + '00000000-0000-0000-0000-000000000000', + 'default_database', + 'default_tenant' + ); \ No newline at end of file diff --git a/go/migrations/20240327075032.sql b/go/migrations/20240327075032.sql new file mode 100644 index 00000000000..e98caf50adb --- /dev/null +++ b/go/migrations/20240327075032.sql @@ -0,0 +1,4 @@ +-- Modify "collections" table +ALTER TABLE "collections" DROP COLUMN "topic"; +-- Modify "segments" table +ALTER TABLE "segments" DROP COLUMN "topic"; diff --git a/go/migrations/20240327172649.sql b/go/migrations/20240327172649.sql new file mode 100644 index 00000000000..0af6002b5af --- /dev/null +++ b/go/migrations/20240327172649.sql @@ -0,0 +1,4 @@ +-- Modify "record_logs" table +-- NOTE: This is a destructive migration autogenerated by atlas. +--This is fine for now because we are still in development. +ALTER TABLE "record_logs" DROP CONSTRAINT "record_logs_pkey", DROP COLUMN "id", ADD COLUMN "log_offset" bigint NOT NULL, ADD PRIMARY KEY ("collection_id", "log_offset"); diff --git a/go/migrations/atlas.sum b/go/migrations/atlas.sum new file mode 100644 index 00000000000..6361b871265 --- /dev/null +++ b/go/migrations/atlas.sum @@ -0,0 +1,5 @@ +h1:9rYxc6RcMJ3Cd4SPZoQ+T6XAUZ7yyWCEnwoRkj1af3c= +20240313233558.sql h1:Gv0TiSYsqGoOZ2T2IWvX4BOasauxool8PrBOIjmmIdg= +20240321194713.sql h1:kVkNpqSFhrXGVGFFvL7JdK3Bw31twFcEhI6A0oCFCkg= +20240327075032.sql h1:nlr2J74XRU8erzHnKJgMr/tKqJxw9+R6RiiEBuvuzgo= +20240327172649.sql h1:UUGo6AzWXKLcpYVd5qH6Hv9jpHNV86z42o6ft5OR0zU= diff --git a/go/mocks/Catalog.go b/go/mocks/Catalog.go new file mode 100644 index 00000000000..d7ce72ebec5 --- /dev/null +++ b/go/mocks/Catalog.go @@ -0,0 +1,526 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + + mock "github.com/stretchr/testify/mock" + + model "github.com/chroma-core/chroma/go/pkg/model" + + types "github.com/chroma-core/chroma/go/pkg/types" +) + +// Catalog is an autogenerated mock type for the Catalog type +type Catalog struct { + mock.Mock +} + +// CreateCollection provides a mock function with given fields: ctx, createCollection, ts +func (_m *Catalog) CreateCollection(ctx context.Context, createCollection *model.CreateCollection, ts int64) (*model.Collection, error) { + ret := _m.Called(ctx, createCollection, ts) + + if len(ret) == 0 { + panic("no return value specified for CreateCollection") + } + + var r0 *model.Collection + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateCollection, int64) (*model.Collection, error)); ok { + return rf(ctx, createCollection, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateCollection, int64) *model.Collection); ok { + r0 = rf(ctx, createCollection, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Collection) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateCollection, int64) error); ok { + r1 = rf(ctx, createCollection, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateDatabase provides a mock function with given fields: ctx, createDatabase, ts +func (_m *Catalog) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase, ts int64) (*model.Database, error) { + ret := _m.Called(ctx, createDatabase, ts) + + if len(ret) == 0 { + panic("no return value specified for CreateDatabase") + } + + var r0 *model.Database + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateDatabase, int64) (*model.Database, error)); ok { + return rf(ctx, createDatabase, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateDatabase, int64) *model.Database); ok { + r0 = rf(ctx, createDatabase, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Database) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateDatabase, int64) error); ok { + r1 = rf(ctx, createDatabase, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateSegment provides a mock function with given fields: ctx, createSegment, ts +func (_m *Catalog) CreateSegment(ctx context.Context, createSegment *model.CreateSegment, ts int64) (*model.Segment, error) { + ret := _m.Called(ctx, createSegment, ts) + + if len(ret) == 0 { + panic("no return value specified for CreateSegment") + } + + var r0 *model.Segment + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateSegment, int64) (*model.Segment, error)); ok { + return rf(ctx, createSegment, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateSegment, int64) *model.Segment); ok { + r0 = rf(ctx, createSegment, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Segment) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateSegment, int64) error); ok { + r1 = rf(ctx, createSegment, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateTenant provides a mock function with given fields: ctx, createTenant, ts +func (_m *Catalog) CreateTenant(ctx context.Context, createTenant *model.CreateTenant, ts int64) (*model.Tenant, error) { + ret := _m.Called(ctx, createTenant, ts) + + if len(ret) == 0 { + panic("no return value specified for CreateTenant") + } + + var r0 *model.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateTenant, int64) (*model.Tenant, error)); ok { + return rf(ctx, createTenant, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateTenant, int64) *model.Tenant); ok { + r0 = rf(ctx, createTenant, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateTenant, int64) error); ok { + r1 = rf(ctx, createTenant, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// DeleteCollection provides a mock function with given fields: ctx, deleteCollection +func (_m *Catalog) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error { + ret := _m.Called(ctx, deleteCollection) + + if len(ret) == 0 { + panic("no return value specified for DeleteCollection") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, *model.DeleteCollection) error); ok { + r0 = rf(ctx, deleteCollection) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteSegment provides a mock function with given fields: ctx, segmentID +func (_m *Catalog) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error { + ret := _m.Called(ctx, segmentID) + + if len(ret) == 0 { + panic("no return value specified for DeleteSegment") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID) error); ok { + r0 = rf(ctx, segmentID) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// FlushCollectionCompaction provides a mock function with given fields: ctx, flushCollectionCompaction +func (_m *Catalog) FlushCollectionCompaction(ctx context.Context, flushCollectionCompaction *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error) { + ret := _m.Called(ctx, flushCollectionCompaction) + + if len(ret) == 0 { + panic("no return value specified for FlushCollectionCompaction") + } + + var r0 *model.FlushCollectionInfo + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error)); ok { + return rf(ctx, flushCollectionCompaction) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.FlushCollectionCompaction) *model.FlushCollectionInfo); ok { + r0 = rf(ctx, flushCollectionCompaction) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.FlushCollectionInfo) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.FlushCollectionCompaction) error); ok { + r1 = rf(ctx, flushCollectionCompaction) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetAllDatabases provides a mock function with given fields: ctx, ts +func (_m *Catalog) GetAllDatabases(ctx context.Context, ts int64) ([]*model.Database, error) { + ret := _m.Called(ctx, ts) + + if len(ret) == 0 { + panic("no return value specified for GetAllDatabases") + } + + var r0 []*model.Database + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, int64) ([]*model.Database, error)); ok { + return rf(ctx, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, int64) []*model.Database); ok { + r0 = rf(ctx, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*model.Database) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok { + r1 = rf(ctx, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetAllTenants provides a mock function with given fields: ctx, ts +func (_m *Catalog) GetAllTenants(ctx context.Context, ts int64) ([]*model.Tenant, error) { + ret := _m.Called(ctx, ts) + + if len(ret) == 0 { + panic("no return value specified for GetAllTenants") + } + + var r0 []*model.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, int64) ([]*model.Tenant, error)); ok { + return rf(ctx, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, int64) []*model.Tenant); ok { + r0 = rf(ctx, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*model.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok { + r1 = rf(ctx, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetCollections provides a mock function with given fields: ctx, collectionID, collectionName, tenantID, databaseName +func (_m *Catalog) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, tenantID string, databaseName string) ([]*model.Collection, error) { + ret := _m.Called(ctx, collectionID, collectionName, tenantID, databaseName) + + if len(ret) == 0 { + panic("no return value specified for GetCollections") + } + + var r0 []*model.Collection + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, string, string) ([]*model.Collection, error)); ok { + return rf(ctx, collectionID, collectionName, tenantID, databaseName) + } + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, string, string) []*model.Collection); ok { + r0 = rf(ctx, collectionID, collectionName, tenantID, databaseName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*model.Collection) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, string, string) error); ok { + r1 = rf(ctx, collectionID, collectionName, tenantID, databaseName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetDatabases provides a mock function with given fields: ctx, getDatabase, ts +func (_m *Catalog) GetDatabases(ctx context.Context, getDatabase *model.GetDatabase, ts int64) (*model.Database, error) { + ret := _m.Called(ctx, getDatabase, ts) + + if len(ret) == 0 { + panic("no return value specified for GetDatabases") + } + + var r0 *model.Database + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.GetDatabase, int64) (*model.Database, error)); ok { + return rf(ctx, getDatabase, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.GetDatabase, int64) *model.Database); ok { + r0 = rf(ctx, getDatabase, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Database) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.GetDatabase, int64) error); ok { + r1 = rf(ctx, getDatabase, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetSegments provides a mock function with given fields: ctx, segmentID, segmentType, scope, collectionID +func (_m *Catalog) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*model.Segment, error) { + ret := _m.Called(ctx, segmentID, segmentType, scope, collectionID) + + if len(ret) == 0 { + panic("no return value specified for GetSegments") + } + + var r0 []*model.Segment + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID) ([]*model.Segment, error)); ok { + return rf(ctx, segmentID, segmentType, scope, collectionID) + } + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID) []*model.Segment); ok { + r0 = rf(ctx, segmentID, segmentType, scope, collectionID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*model.Segment) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID) error); ok { + r1 = rf(ctx, segmentID, segmentType, scope, collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenants provides a mock function with given fields: ctx, getTenant, ts +func (_m *Catalog) GetTenants(ctx context.Context, getTenant *model.GetTenant, ts int64) (*model.Tenant, error) { + ret := _m.Called(ctx, getTenant, ts) + + if len(ret) == 0 { + panic("no return value specified for GetTenants") + } + + var r0 *model.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.GetTenant, int64) (*model.Tenant, error)); ok { + return rf(ctx, getTenant, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.GetTenant, int64) *model.Tenant); ok { + r0 = rf(ctx, getTenant, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.GetTenant, int64) error); ok { + r1 = rf(ctx, getTenant, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenantsLastCompactionTime provides a mock function with given fields: ctx, tenantIDs +func (_m *Catalog) GetTenantsLastCompactionTime(ctx context.Context, tenantIDs []string) ([]*dbmodel.Tenant, error) { + ret := _m.Called(ctx, tenantIDs) + + if len(ret) == 0 { + panic("no return value specified for GetTenantsLastCompactionTime") + } + + var r0 []*dbmodel.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, []string) ([]*dbmodel.Tenant, error)); ok { + return rf(ctx, tenantIDs) + } + if rf, ok := ret.Get(0).(func(context.Context, []string) []*dbmodel.Tenant); ok { + r0 = rf(ctx, tenantIDs) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok { + r1 = rf(ctx, tenantIDs) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// ResetState provides a mock function with given fields: ctx +func (_m *Catalog) ResetState(ctx context.Context) error { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for ResetState") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context) error); ok { + r0 = rf(ctx) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// SetTenantLastCompactionTime provides a mock function with given fields: ctx, tenantID, lastCompactionTime +func (_m *Catalog) SetTenantLastCompactionTime(ctx context.Context, tenantID string, lastCompactionTime int64) error { + ret := _m.Called(ctx, tenantID, lastCompactionTime) + + if len(ret) == 0 { + panic("no return value specified for SetTenantLastCompactionTime") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, string, int64) error); ok { + r0 = rf(ctx, tenantID, lastCompactionTime) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// UpdateCollection provides a mock function with given fields: ctx, updateCollection, ts +func (_m *Catalog) UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection, ts int64) (*model.Collection, error) { + ret := _m.Called(ctx, updateCollection, ts) + + if len(ret) == 0 { + panic("no return value specified for UpdateCollection") + } + + var r0 *model.Collection + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateCollection, int64) (*model.Collection, error)); ok { + return rf(ctx, updateCollection, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateCollection, int64) *model.Collection); ok { + r0 = rf(ctx, updateCollection, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Collection) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.UpdateCollection, int64) error); ok { + r1 = rf(ctx, updateCollection, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// UpdateSegment provides a mock function with given fields: ctx, segmentInfo, ts +func (_m *Catalog) UpdateSegment(ctx context.Context, segmentInfo *model.UpdateSegment, ts int64) (*model.Segment, error) { + ret := _m.Called(ctx, segmentInfo, ts) + + if len(ret) == 0 { + panic("no return value specified for UpdateSegment") + } + + var r0 *model.Segment + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateSegment, int64) (*model.Segment, error)); ok { + return rf(ctx, segmentInfo, ts) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateSegment, int64) *model.Segment); ok { + r0 = rf(ctx, segmentInfo, ts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Segment) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.UpdateSegment, int64) error); ok { + r1 = rf(ctx, segmentInfo, ts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewCatalog creates a new instance of Catalog. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewCatalog(t interface { + mock.TestingT + Cleanup(func()) +}) *Catalog { + mock := &Catalog{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/CollectionMetadataValueType.go b/go/mocks/CollectionMetadataValueType.go new file mode 100644 index 00000000000..2ed88d4b39c --- /dev/null +++ b/go/mocks/CollectionMetadataValueType.go @@ -0,0 +1,50 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + model "github.com/chroma-core/chroma/go/pkg/model" + mock "github.com/stretchr/testify/mock" +) + +// CollectionMetadataValueType is an autogenerated mock type for the CollectionMetadataValueType type +type CollectionMetadataValueType struct { + mock.Mock +} + +// Equals provides a mock function with given fields: other +func (_m *CollectionMetadataValueType) Equals(other model.CollectionMetadataValueType) bool { + ret := _m.Called(other) + + if len(ret) == 0 { + panic("no return value specified for Equals") + } + + var r0 bool + if rf, ok := ret.Get(0).(func(model.CollectionMetadataValueType) bool); ok { + r0 = rf(other) + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// IsCollectionMetadataValueType provides a mock function with given fields: +func (_m *CollectionMetadataValueType) IsCollectionMetadataValueType() { + _m.Called() +} + +// NewCollectionMetadataValueType creates a new instance of CollectionMetadataValueType. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewCollectionMetadataValueType(t interface { + mock.TestingT + Cleanup(func()) +}) *CollectionMetadataValueType { + mock := &CollectionMetadataValueType{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/Component.go b/go/mocks/Component.go new file mode 100644 index 00000000000..101d3661aca --- /dev/null +++ b/go/mocks/Component.go @@ -0,0 +1,60 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// Component is an autogenerated mock type for the Component type +type Component struct { + mock.Mock +} + +// Start provides a mock function with given fields: +func (_m *Component) Start() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Start") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Stop provides a mock function with given fields: +func (_m *Component) Stop() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Stop") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewComponent creates a new instance of Component. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewComponent(t interface { + mock.TestingT + Cleanup(func()) +}) *Component { + mock := &Component{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/GrpcProvider.go b/go/mocks/GrpcProvider.go new file mode 100644 index 00000000000..2859bcd862a --- /dev/null +++ b/go/mocks/GrpcProvider.go @@ -0,0 +1,58 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + grpcutils "github.com/chroma-core/chroma/go/pkg/grpcutils" + mock "github.com/stretchr/testify/mock" + grpc "google.golang.org/grpc" +) + +// GrpcProvider is an autogenerated mock type for the GrpcProvider type +type GrpcProvider struct { + mock.Mock +} + +// StartGrpcServer provides a mock function with given fields: name, grpcConfig, registerFunc +func (_m *GrpcProvider) StartGrpcServer(name string, grpcConfig *grpcutils.GrpcConfig, registerFunc func(grpc.ServiceRegistrar)) (grpcutils.GrpcServer, error) { + ret := _m.Called(name, grpcConfig, registerFunc) + + if len(ret) == 0 { + panic("no return value specified for StartGrpcServer") + } + + var r0 grpcutils.GrpcServer + var r1 error + if rf, ok := ret.Get(0).(func(string, *grpcutils.GrpcConfig, func(grpc.ServiceRegistrar)) (grpcutils.GrpcServer, error)); ok { + return rf(name, grpcConfig, registerFunc) + } + if rf, ok := ret.Get(0).(func(string, *grpcutils.GrpcConfig, func(grpc.ServiceRegistrar)) grpcutils.GrpcServer); ok { + r0 = rf(name, grpcConfig, registerFunc) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(grpcutils.GrpcServer) + } + } + + if rf, ok := ret.Get(1).(func(string, *grpcutils.GrpcConfig, func(grpc.ServiceRegistrar)) error); ok { + r1 = rf(name, grpcConfig, registerFunc) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewGrpcProvider creates a new instance of GrpcProvider. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewGrpcProvider(t interface { + mock.TestingT + Cleanup(func()) +}) *GrpcProvider { + mock := &GrpcProvider{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/GrpcServer.go b/go/mocks/GrpcServer.go new file mode 100644 index 00000000000..bf55ae5c116 --- /dev/null +++ b/go/mocks/GrpcServer.go @@ -0,0 +1,60 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// GrpcServer is an autogenerated mock type for the GrpcServer type +type GrpcServer struct { + mock.Mock +} + +// Close provides a mock function with given fields: +func (_m *GrpcServer) Close() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Close") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Port provides a mock function with given fields: +func (_m *GrpcServer) Port() int { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Port") + } + + var r0 int + if rf, ok := ret.Get(0).(func() int); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(int) + } + + return r0 +} + +// NewGrpcServer creates a new instance of GrpcServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewGrpcServer(t interface { + mock.TestingT + Cleanup(func()) +}) *GrpcServer { + mock := &GrpcServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/ICollectionDb.go b/go/mocks/ICollectionDb.go new file mode 100644 index 00000000000..889c99781e4 --- /dev/null +++ b/go/mocks/ICollectionDb.go @@ -0,0 +1,167 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// ICollectionDb is an autogenerated mock type for the ICollectionDb type +type ICollectionDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *ICollectionDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteCollectionByID provides a mock function with given fields: collectionID +func (_m *ICollectionDb) DeleteCollectionByID(collectionID string) (int, error) { + ret := _m.Called(collectionID) + + if len(ret) == 0 { + panic("no return value specified for DeleteCollectionByID") + } + + var r0 int + var r1 error + if rf, ok := ret.Get(0).(func(string) (int, error)); ok { + return rf(collectionID) + } + if rf, ok := ret.Get(0).(func(string) int); ok { + r0 = rf(collectionID) + } else { + r0 = ret.Get(0).(int) + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetCollections provides a mock function with given fields: collectionID, collectionName, tenantID, databaseName +func (_m *ICollectionDb) GetCollections(collectionID *string, collectionName *string, tenantID string, databaseName string) ([]*dbmodel.CollectionAndMetadata, error) { + ret := _m.Called(collectionID, collectionName, tenantID, databaseName) + + if len(ret) == 0 { + panic("no return value specified for GetCollections") + } + + var r0 []*dbmodel.CollectionAndMetadata + var r1 error + if rf, ok := ret.Get(0).(func(*string, *string, string, string) ([]*dbmodel.CollectionAndMetadata, error)); ok { + return rf(collectionID, collectionName, tenantID, databaseName) + } + if rf, ok := ret.Get(0).(func(*string, *string, string, string) []*dbmodel.CollectionAndMetadata); ok { + r0 = rf(collectionID, collectionName, tenantID, databaseName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.CollectionAndMetadata) + } + } + + if rf, ok := ret.Get(1).(func(*string, *string, string, string) error); ok { + r1 = rf(collectionID, collectionName, tenantID, databaseName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: in +func (_m *ICollectionDb) Insert(in *dbmodel.Collection) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Collection) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Update provides a mock function with given fields: in +func (_m *ICollectionDb) Update(in *dbmodel.Collection) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Update") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Collection) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// UpdateLogPositionAndVersion provides a mock function with given fields: collectionID, logPosition, currentCollectionVersion +func (_m *ICollectionDb) UpdateLogPositionAndVersion(collectionID string, logPosition int64, currentCollectionVersion int32) (int32, error) { + ret := _m.Called(collectionID, logPosition, currentCollectionVersion) + + if len(ret) == 0 { + panic("no return value specified for UpdateLogPositionAndVersion") + } + + var r0 int32 + var r1 error + if rf, ok := ret.Get(0).(func(string, int64, int32) (int32, error)); ok { + return rf(collectionID, logPosition, currentCollectionVersion) + } + if rf, ok := ret.Get(0).(func(string, int64, int32) int32); ok { + r0 = rf(collectionID, logPosition, currentCollectionVersion) + } else { + r0 = ret.Get(0).(int32) + } + + if rf, ok := ret.Get(1).(func(string, int64, int32) error); ok { + r1 = rf(collectionID, logPosition, currentCollectionVersion) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewICollectionDb creates a new instance of ICollectionDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewICollectionDb(t interface { + mock.TestingT + Cleanup(func()) +}) *ICollectionDb { + mock := &ICollectionDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/ICollectionMetadataDb.go b/go/mocks/ICollectionMetadataDb.go similarity index 67% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/ICollectionMetadataDb.go rename to go/mocks/ICollectionMetadataDb.go index 87d71909b06..d231bf1fe29 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/ICollectionMetadataDb.go +++ b/go/mocks/ICollectionMetadataDb.go @@ -1,9 +1,9 @@ -// Code generated by mockery v2.33.3. DO NOT EDIT. +// Code generated by mockery v2.42.1. DO NOT EDIT. package mocks import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" ) @@ -16,6 +16,10 @@ type ICollectionMetadataDb struct { func (_m *ICollectionMetadataDb) DeleteAll() error { ret := _m.Called() + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + var r0 error if rf, ok := ret.Get(0).(func() error); ok { r0 = rf() @@ -27,23 +31,41 @@ func (_m *ICollectionMetadataDb) DeleteAll() error { } // DeleteByCollectionID provides a mock function with given fields: collectionID -func (_m *ICollectionMetadataDb) DeleteByCollectionID(collectionID string) error { +func (_m *ICollectionMetadataDb) DeleteByCollectionID(collectionID string) (int, error) { ret := _m.Called(collectionID) - var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { + if len(ret) == 0 { + panic("no return value specified for DeleteByCollectionID") + } + + var r0 int + var r1 error + if rf, ok := ret.Get(0).(func(string) (int, error)); ok { + return rf(collectionID) + } + if rf, ok := ret.Get(0).(func(string) int); ok { r0 = rf(collectionID) } else { - r0 = ret.Error(0) + r0 = ret.Get(0).(int) } - return r0 + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 } // Insert provides a mock function with given fields: in func (_m *ICollectionMetadataDb) Insert(in []*dbmodel.CollectionMetadata) error { ret := _m.Called(in) + if len(ret) == 0 { + panic("no return value specified for Insert") + } + var r0 error if rf, ok := ret.Get(0).(func([]*dbmodel.CollectionMetadata) error); ok { r0 = rf(in) diff --git a/go/mocks/ICoordinator.go b/go/mocks/ICoordinator.go new file mode 100644 index 00000000000..02d51811196 --- /dev/null +++ b/go/mocks/ICoordinator.go @@ -0,0 +1,490 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + + mock "github.com/stretchr/testify/mock" + + model "github.com/chroma-core/chroma/go/pkg/model" + + types "github.com/chroma-core/chroma/go/pkg/types" +) + +// ICoordinator is an autogenerated mock type for the ICoordinator type +type ICoordinator struct { + mock.Mock +} + +// CreateCollection provides a mock function with given fields: ctx, createCollection +func (_m *ICoordinator) CreateCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error) { + ret := _m.Called(ctx, createCollection) + + if len(ret) == 0 { + panic("no return value specified for CreateCollection") + } + + var r0 *model.Collection + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateCollection) (*model.Collection, error)); ok { + return rf(ctx, createCollection) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateCollection) *model.Collection); ok { + r0 = rf(ctx, createCollection) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Collection) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateCollection) error); ok { + r1 = rf(ctx, createCollection) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateDatabase provides a mock function with given fields: ctx, createDatabase +func (_m *ICoordinator) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error) { + ret := _m.Called(ctx, createDatabase) + + if len(ret) == 0 { + panic("no return value specified for CreateDatabase") + } + + var r0 *model.Database + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateDatabase) (*model.Database, error)); ok { + return rf(ctx, createDatabase) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateDatabase) *model.Database); ok { + r0 = rf(ctx, createDatabase) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Database) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateDatabase) error); ok { + r1 = rf(ctx, createDatabase) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateSegment provides a mock function with given fields: ctx, createSegment +func (_m *ICoordinator) CreateSegment(ctx context.Context, createSegment *model.CreateSegment) error { + ret := _m.Called(ctx, createSegment) + + if len(ret) == 0 { + panic("no return value specified for CreateSegment") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateSegment) error); ok { + r0 = rf(ctx, createSegment) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// CreateTenant provides a mock function with given fields: ctx, createTenant +func (_m *ICoordinator) CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) { + ret := _m.Called(ctx, createTenant) + + if len(ret) == 0 { + panic("no return value specified for CreateTenant") + } + + var r0 *model.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateTenant) (*model.Tenant, error)); ok { + return rf(ctx, createTenant) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.CreateTenant) *model.Tenant); ok { + r0 = rf(ctx, createTenant) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.CreateTenant) error); ok { + r1 = rf(ctx, createTenant) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// DeleteCollection provides a mock function with given fields: ctx, deleteCollection +func (_m *ICoordinator) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error { + ret := _m.Called(ctx, deleteCollection) + + if len(ret) == 0 { + panic("no return value specified for DeleteCollection") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, *model.DeleteCollection) error); ok { + r0 = rf(ctx, deleteCollection) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteSegment provides a mock function with given fields: ctx, segmentID +func (_m *ICoordinator) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error { + ret := _m.Called(ctx, segmentID) + + if len(ret) == 0 { + panic("no return value specified for DeleteSegment") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID) error); ok { + r0 = rf(ctx, segmentID) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// FlushCollectionCompaction provides a mock function with given fields: ctx, flushCollectionCompaction +func (_m *ICoordinator) FlushCollectionCompaction(ctx context.Context, flushCollectionCompaction *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error) { + ret := _m.Called(ctx, flushCollectionCompaction) + + if len(ret) == 0 { + panic("no return value specified for FlushCollectionCompaction") + } + + var r0 *model.FlushCollectionInfo + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error)); ok { + return rf(ctx, flushCollectionCompaction) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.FlushCollectionCompaction) *model.FlushCollectionInfo); ok { + r0 = rf(ctx, flushCollectionCompaction) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.FlushCollectionInfo) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.FlushCollectionCompaction) error); ok { + r1 = rf(ctx, flushCollectionCompaction) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetCollections provides a mock function with given fields: ctx, collectionID, collectionName, tenantID, dataName +func (_m *ICoordinator) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, tenantID string, dataName string) ([]*model.Collection, error) { + ret := _m.Called(ctx, collectionID, collectionName, tenantID, dataName) + + if len(ret) == 0 { + panic("no return value specified for GetCollections") + } + + var r0 []*model.Collection + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, string, string) ([]*model.Collection, error)); ok { + return rf(ctx, collectionID, collectionName, tenantID, dataName) + } + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, string, string) []*model.Collection); ok { + r0 = rf(ctx, collectionID, collectionName, tenantID, dataName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*model.Collection) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, string, string) error); ok { + r1 = rf(ctx, collectionID, collectionName, tenantID, dataName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetDatabase provides a mock function with given fields: ctx, getDatabase +func (_m *ICoordinator) GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error) { + ret := _m.Called(ctx, getDatabase) + + if len(ret) == 0 { + panic("no return value specified for GetDatabase") + } + + var r0 *model.Database + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.GetDatabase) (*model.Database, error)); ok { + return rf(ctx, getDatabase) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.GetDatabase) *model.Database); ok { + r0 = rf(ctx, getDatabase) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Database) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.GetDatabase) error); ok { + r1 = rf(ctx, getDatabase) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetSegments provides a mock function with given fields: ctx, segmentID, segmentType, scope, collectionID +func (_m *ICoordinator) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*model.Segment, error) { + ret := _m.Called(ctx, segmentID, segmentType, scope, collectionID) + + if len(ret) == 0 { + panic("no return value specified for GetSegments") + } + + var r0 []*model.Segment + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID) ([]*model.Segment, error)); ok { + return rf(ctx, segmentID, segmentType, scope, collectionID) + } + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID) []*model.Segment); ok { + r0 = rf(ctx, segmentID, segmentType, scope, collectionID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*model.Segment) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID) error); ok { + r1 = rf(ctx, segmentID, segmentType, scope, collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenant provides a mock function with given fields: ctx, getTenant +func (_m *ICoordinator) GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error) { + ret := _m.Called(ctx, getTenant) + + if len(ret) == 0 { + panic("no return value specified for GetTenant") + } + + var r0 *model.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.GetTenant) (*model.Tenant, error)); ok { + return rf(ctx, getTenant) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.GetTenant) *model.Tenant); ok { + r0 = rf(ctx, getTenant) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.GetTenant) error); ok { + r1 = rf(ctx, getTenant) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenantsLastCompactionTime provides a mock function with given fields: ctx, tenantIDs +func (_m *ICoordinator) GetTenantsLastCompactionTime(ctx context.Context, tenantIDs []string) ([]*dbmodel.Tenant, error) { + ret := _m.Called(ctx, tenantIDs) + + if len(ret) == 0 { + panic("no return value specified for GetTenantsLastCompactionTime") + } + + var r0 []*dbmodel.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, []string) ([]*dbmodel.Tenant, error)); ok { + return rf(ctx, tenantIDs) + } + if rf, ok := ret.Get(0).(func(context.Context, []string) []*dbmodel.Tenant); ok { + r0 = rf(ctx, tenantIDs) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok { + r1 = rf(ctx, tenantIDs) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// ResetState provides a mock function with given fields: ctx +func (_m *ICoordinator) ResetState(ctx context.Context) error { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for ResetState") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context) error); ok { + r0 = rf(ctx) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// SetTenantLastCompactionTime provides a mock function with given fields: ctx, tenantID, lastCompactionTime +func (_m *ICoordinator) SetTenantLastCompactionTime(ctx context.Context, tenantID string, lastCompactionTime int64) error { + ret := _m.Called(ctx, tenantID, lastCompactionTime) + + if len(ret) == 0 { + panic("no return value specified for SetTenantLastCompactionTime") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, string, int64) error); ok { + r0 = rf(ctx, tenantID, lastCompactionTime) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Start provides a mock function with given fields: +func (_m *ICoordinator) Start() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Start") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Stop provides a mock function with given fields: +func (_m *ICoordinator) Stop() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Stop") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// UpdateCollection provides a mock function with given fields: ctx, updateCollection +func (_m *ICoordinator) UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection) (*model.Collection, error) { + ret := _m.Called(ctx, updateCollection) + + if len(ret) == 0 { + panic("no return value specified for UpdateCollection") + } + + var r0 *model.Collection + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateCollection) (*model.Collection, error)); ok { + return rf(ctx, updateCollection) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateCollection) *model.Collection); ok { + r0 = rf(ctx, updateCollection) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Collection) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.UpdateCollection) error); ok { + r1 = rf(ctx, updateCollection) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// UpdateSegment provides a mock function with given fields: ctx, updateSegment +func (_m *ICoordinator) UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment) (*model.Segment, error) { + ret := _m.Called(ctx, updateSegment) + + if len(ret) == 0 { + panic("no return value specified for UpdateSegment") + } + + var r0 *model.Segment + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateSegment) (*model.Segment, error)); ok { + return rf(ctx, updateSegment) + } + if rf, ok := ret.Get(0).(func(context.Context, *model.UpdateSegment) *model.Segment); ok { + r0 = rf(ctx, updateSegment) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Segment) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *model.UpdateSegment) error); ok { + r1 = rf(ctx, updateSegment) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewICoordinator creates a new instance of ICoordinator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewICoordinator(t interface { + mock.TestingT + Cleanup(func()) +}) *ICoordinator { + mock := &ICoordinator{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IDatabaseDb.go b/go/mocks/IDatabaseDb.go new file mode 100644 index 00000000000..982402e946d --- /dev/null +++ b/go/mocks/IDatabaseDb.go @@ -0,0 +1,123 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// IDatabaseDb is an autogenerated mock type for the IDatabaseDb type +type IDatabaseDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *IDatabaseDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// GetAllDatabases provides a mock function with given fields: +func (_m *IDatabaseDb) GetAllDatabases() ([]*dbmodel.Database, error) { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for GetAllDatabases") + } + + var r0 []*dbmodel.Database + var r1 error + if rf, ok := ret.Get(0).(func() ([]*dbmodel.Database, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []*dbmodel.Database); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Database) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetDatabases provides a mock function with given fields: tenantID, databaseName +func (_m *IDatabaseDb) GetDatabases(tenantID string, databaseName string) ([]*dbmodel.Database, error) { + ret := _m.Called(tenantID, databaseName) + + if len(ret) == 0 { + panic("no return value specified for GetDatabases") + } + + var r0 []*dbmodel.Database + var r1 error + if rf, ok := ret.Get(0).(func(string, string) ([]*dbmodel.Database, error)); ok { + return rf(tenantID, databaseName) + } + if rf, ok := ret.Get(0).(func(string, string) []*dbmodel.Database); ok { + r0 = rf(tenantID, databaseName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Database) + } + } + + if rf, ok := ret.Get(1).(func(string, string) error); ok { + r1 = rf(tenantID, databaseName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: in +func (_m *IDatabaseDb) Insert(in *dbmodel.Database) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Database) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewIDatabaseDb creates a new instance of IDatabaseDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIDatabaseDb(t interface { + mock.TestingT + Cleanup(func()) +}) *IDatabaseDb { + mock := &IDatabaseDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IMemberlistManager.go b/go/mocks/IMemberlistManager.go new file mode 100644 index 00000000000..a5bd93a54a1 --- /dev/null +++ b/go/mocks/IMemberlistManager.go @@ -0,0 +1,60 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// IMemberlistManager is an autogenerated mock type for the IMemberlistManager type +type IMemberlistManager struct { + mock.Mock +} + +// Start provides a mock function with given fields: +func (_m *IMemberlistManager) Start() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Start") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Stop provides a mock function with given fields: +func (_m *IMemberlistManager) Stop() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Stop") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewIMemberlistManager creates a new instance of IMemberlistManager. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIMemberlistManager(t interface { + mock.TestingT + Cleanup(func()) +}) *IMemberlistManager { + mock := &IMemberlistManager{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IMemberlistStore.go b/go/mocks/IMemberlistStore.go new file mode 100644 index 00000000000..06262ece9b2 --- /dev/null +++ b/go/mocks/IMemberlistStore.go @@ -0,0 +1,84 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + memberlist_manager "github.com/chroma-core/chroma/go/pkg/memberlist_manager" + mock "github.com/stretchr/testify/mock" +) + +// IMemberlistStore is an autogenerated mock type for the IMemberlistStore type +type IMemberlistStore struct { + mock.Mock +} + +// GetMemberlist provides a mock function with given fields: ctx +func (_m *IMemberlistStore) GetMemberlist(ctx context.Context) (*memberlist_manager.Memberlist, string, error) { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for GetMemberlist") + } + + var r0 *memberlist_manager.Memberlist + var r1 string + var r2 error + if rf, ok := ret.Get(0).(func(context.Context) (*memberlist_manager.Memberlist, string, error)); ok { + return rf(ctx) + } + if rf, ok := ret.Get(0).(func(context.Context) *memberlist_manager.Memberlist); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*memberlist_manager.Memberlist) + } + } + + if rf, ok := ret.Get(1).(func(context.Context) string); ok { + r1 = rf(ctx) + } else { + r1 = ret.Get(1).(string) + } + + if rf, ok := ret.Get(2).(func(context.Context) error); ok { + r2 = rf(ctx) + } else { + r2 = ret.Error(2) + } + + return r0, r1, r2 +} + +// UpdateMemberlist provides a mock function with given fields: ctx, memberlist, resourceVersion +func (_m *IMemberlistStore) UpdateMemberlist(ctx context.Context, memberlist *memberlist_manager.Memberlist, resourceVersion string) error { + ret := _m.Called(ctx, memberlist, resourceVersion) + + if len(ret) == 0 { + panic("no return value specified for UpdateMemberlist") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, *memberlist_manager.Memberlist, string) error); ok { + r0 = rf(ctx, memberlist, resourceVersion) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewIMemberlistStore creates a new instance of IMemberlistStore. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIMemberlistStore(t interface { + mock.TestingT + Cleanup(func()) +}) *IMemberlistStore { + mock := &IMemberlistStore{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IMetaDomain.go b/go/mocks/IMetaDomain.go new file mode 100644 index 00000000000..e4b4bb130d5 --- /dev/null +++ b/go/mocks/IMetaDomain.go @@ -0,0 +1,189 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// IMetaDomain is an autogenerated mock type for the IMetaDomain type +type IMetaDomain struct { + mock.Mock +} + +// CollectionDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) CollectionDb(ctx context.Context) dbmodel.ICollectionDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for CollectionDb") + } + + var r0 dbmodel.ICollectionDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.ICollectionDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.ICollectionDb) + } + } + + return r0 +} + +// CollectionMetadataDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) CollectionMetadataDb(ctx context.Context) dbmodel.ICollectionMetadataDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for CollectionMetadataDb") + } + + var r0 dbmodel.ICollectionMetadataDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.ICollectionMetadataDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.ICollectionMetadataDb) + } + } + + return r0 +} + +// DatabaseDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) DatabaseDb(ctx context.Context) dbmodel.IDatabaseDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for DatabaseDb") + } + + var r0 dbmodel.IDatabaseDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.IDatabaseDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.IDatabaseDb) + } + } + + return r0 +} + +// NotificationDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) NotificationDb(ctx context.Context) dbmodel.INotificationDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for NotificationDb") + } + + var r0 dbmodel.INotificationDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.INotificationDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.INotificationDb) + } + } + + return r0 +} + +// RecordLogDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) RecordLogDb(ctx context.Context) dbmodel.IRecordLogDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for RecordLogDb") + } + + var r0 dbmodel.IRecordLogDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.IRecordLogDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.IRecordLogDb) + } + } + + return r0 +} + +// SegmentDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) SegmentDb(ctx context.Context) dbmodel.ISegmentDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for SegmentDb") + } + + var r0 dbmodel.ISegmentDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.ISegmentDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.ISegmentDb) + } + } + + return r0 +} + +// SegmentMetadataDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) SegmentMetadataDb(ctx context.Context) dbmodel.ISegmentMetadataDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for SegmentMetadataDb") + } + + var r0 dbmodel.ISegmentMetadataDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.ISegmentMetadataDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.ISegmentMetadataDb) + } + } + + return r0 +} + +// TenantDb provides a mock function with given fields: ctx +func (_m *IMetaDomain) TenantDb(ctx context.Context) dbmodel.ITenantDb { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for TenantDb") + } + + var r0 dbmodel.ITenantDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.ITenantDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.ITenantDb) + } + } + + return r0 +} + +// NewIMetaDomain creates a new instance of IMetaDomain. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIMetaDomain(t interface { + mock.TestingT + Cleanup(func()) +}) *IMetaDomain { + mock := &IMetaDomain{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/INotificationDb.go b/go/mocks/INotificationDb.go new file mode 100644 index 00000000000..3a0a3d019cf --- /dev/null +++ b/go/mocks/INotificationDb.go @@ -0,0 +1,141 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// INotificationDb is an autogenerated mock type for the INotificationDb type +type INotificationDb struct { + mock.Mock +} + +// Delete provides a mock function with given fields: id +func (_m *INotificationDb) Delete(id []int64) error { + ret := _m.Called(id) + + if len(ret) == 0 { + panic("no return value specified for Delete") + } + + var r0 error + if rf, ok := ret.Get(0).(func([]int64) error); ok { + r0 = rf(id) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteAll provides a mock function with given fields: +func (_m *INotificationDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// GetAllPendingNotifications provides a mock function with given fields: +func (_m *INotificationDb) GetAllPendingNotifications() ([]*dbmodel.Notification, error) { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for GetAllPendingNotifications") + } + + var r0 []*dbmodel.Notification + var r1 error + if rf, ok := ret.Get(0).(func() ([]*dbmodel.Notification, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []*dbmodel.Notification); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Notification) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetNotificationByCollectionID provides a mock function with given fields: collectionID +func (_m *INotificationDb) GetNotificationByCollectionID(collectionID string) ([]*dbmodel.Notification, error) { + ret := _m.Called(collectionID) + + if len(ret) == 0 { + panic("no return value specified for GetNotificationByCollectionID") + } + + var r0 []*dbmodel.Notification + var r1 error + if rf, ok := ret.Get(0).(func(string) ([]*dbmodel.Notification, error)); ok { + return rf(collectionID) + } + if rf, ok := ret.Get(0).(func(string) []*dbmodel.Notification); ok { + r0 = rf(collectionID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Notification) + } + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: in +func (_m *INotificationDb) Insert(in *dbmodel.Notification) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Notification) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewINotificationDb creates a new instance of INotificationDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewINotificationDb(t interface { + mock.TestingT + Cleanup(func()) +}) *INotificationDb { + mock := &INotificationDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IRecordLog.go b/go/mocks/IRecordLog.go new file mode 100644 index 00000000000..885ea7f35ae --- /dev/null +++ b/go/mocks/IRecordLog.go @@ -0,0 +1,156 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + + mock "github.com/stretchr/testify/mock" + + types "github.com/chroma-core/chroma/go/pkg/types" +) + +// IRecordLog is an autogenerated mock type for the IRecordLog type +type IRecordLog struct { + mock.Mock +} + +// GetAllCollectionIDsToCompact provides a mock function with given fields: +func (_m *IRecordLog) GetAllCollectionIDsToCompact() ([]*dbmodel.RecordLog, error) { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for GetAllCollectionIDsToCompact") + } + + var r0 []*dbmodel.RecordLog + var r1 error + if rf, ok := ret.Get(0).(func() ([]*dbmodel.RecordLog, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []*dbmodel.RecordLog); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.RecordLog) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PullLogs provides a mock function with given fields: ctx, collectionID, id, batchSize, endTimestamp +func (_m *IRecordLog) PullLogs(ctx context.Context, collectionID types.UniqueID, id int64, batchSize int, endTimestamp int64) ([]*dbmodel.RecordLog, error) { + ret := _m.Called(ctx, collectionID, id, batchSize, endTimestamp) + + if len(ret) == 0 { + panic("no return value specified for PullLogs") + } + + var r0 []*dbmodel.RecordLog + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, int64, int, int64) ([]*dbmodel.RecordLog, error)); ok { + return rf(ctx, collectionID, id, batchSize, endTimestamp) + } + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, int64, int, int64) []*dbmodel.RecordLog); ok { + r0 = rf(ctx, collectionID, id, batchSize, endTimestamp) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.RecordLog) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, int64, int, int64) error); ok { + r1 = rf(ctx, collectionID, id, batchSize, endTimestamp) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PushLogs provides a mock function with given fields: ctx, collectionID, recordContent +func (_m *IRecordLog) PushLogs(ctx context.Context, collectionID types.UniqueID, recordContent [][]byte) (int, error) { + ret := _m.Called(ctx, collectionID, recordContent) + + if len(ret) == 0 { + panic("no return value specified for PushLogs") + } + + var r0 int + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, [][]byte) (int, error)); ok { + return rf(ctx, collectionID, recordContent) + } + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, [][]byte) int); ok { + r0 = rf(ctx, collectionID, recordContent) + } else { + r0 = ret.Get(0).(int) + } + + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, [][]byte) error); ok { + r1 = rf(ctx, collectionID, recordContent) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Start provides a mock function with given fields: +func (_m *IRecordLog) Start() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Start") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Stop provides a mock function with given fields: +func (_m *IRecordLog) Stop() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Stop") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewIRecordLog creates a new instance of IRecordLog. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIRecordLog(t interface { + mock.TestingT + Cleanup(func()) +}) *IRecordLog { + mock := &IRecordLog{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IRecordLogDb.go b/go/mocks/IRecordLogDb.go new file mode 100644 index 00000000000..bd4dee0281e --- /dev/null +++ b/go/mocks/IRecordLogDb.go @@ -0,0 +1,117 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" + + types "github.com/chroma-core/chroma/go/pkg/types" +) + +// IRecordLogDb is an autogenerated mock type for the IRecordLogDb type +type IRecordLogDb struct { + mock.Mock +} + +// GetAllCollectionsToCompact provides a mock function with given fields: +func (_m *IRecordLogDb) GetAllCollectionsToCompact() ([]*dbmodel.RecordLog, error) { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for GetAllCollectionsToCompact") + } + + var r0 []*dbmodel.RecordLog + var r1 error + if rf, ok := ret.Get(0).(func() ([]*dbmodel.RecordLog, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []*dbmodel.RecordLog); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.RecordLog) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PullLogs provides a mock function with given fields: collectionID, id, batchSize, endTimestamp +func (_m *IRecordLogDb) PullLogs(collectionID types.UniqueID, id int64, batchSize int, endTimestamp int64) ([]*dbmodel.RecordLog, error) { + ret := _m.Called(collectionID, id, batchSize, endTimestamp) + + if len(ret) == 0 { + panic("no return value specified for PullLogs") + } + + var r0 []*dbmodel.RecordLog + var r1 error + if rf, ok := ret.Get(0).(func(types.UniqueID, int64, int, int64) ([]*dbmodel.RecordLog, error)); ok { + return rf(collectionID, id, batchSize, endTimestamp) + } + if rf, ok := ret.Get(0).(func(types.UniqueID, int64, int, int64) []*dbmodel.RecordLog); ok { + r0 = rf(collectionID, id, batchSize, endTimestamp) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.RecordLog) + } + } + + if rf, ok := ret.Get(1).(func(types.UniqueID, int64, int, int64) error); ok { + r1 = rf(collectionID, id, batchSize, endTimestamp) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PushLogs provides a mock function with given fields: collectionID, recordsContent +func (_m *IRecordLogDb) PushLogs(collectionID types.UniqueID, recordsContent [][]byte) (int, error) { + ret := _m.Called(collectionID, recordsContent) + + if len(ret) == 0 { + panic("no return value specified for PushLogs") + } + + var r0 int + var r1 error + if rf, ok := ret.Get(0).(func(types.UniqueID, [][]byte) (int, error)); ok { + return rf(collectionID, recordsContent) + } + if rf, ok := ret.Get(0).(func(types.UniqueID, [][]byte) int); ok { + r0 = rf(collectionID, recordsContent) + } else { + r0 = ret.Get(0).(int) + } + + if rf, ok := ret.Get(1).(func(types.UniqueID, [][]byte) error); ok { + r1 = rf(collectionID, recordsContent) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewIRecordLogDb creates a new instance of IRecordLogDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIRecordLogDb(t interface { + mock.TestingT + Cleanup(func()) +}) *IRecordLogDb { + mock := &IRecordLogDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/ISegmentDb.go b/go/mocks/ISegmentDb.go new file mode 100644 index 00000000000..06d1f104090 --- /dev/null +++ b/go/mocks/ISegmentDb.go @@ -0,0 +1,151 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" + + model "github.com/chroma-core/chroma/go/pkg/model" + + types "github.com/chroma-core/chroma/go/pkg/types" +) + +// ISegmentDb is an autogenerated mock type for the ISegmentDb type +type ISegmentDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *ISegmentDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteSegmentByID provides a mock function with given fields: id +func (_m *ISegmentDb) DeleteSegmentByID(id string) error { + ret := _m.Called(id) + + if len(ret) == 0 { + panic("no return value specified for DeleteSegmentByID") + } + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(id) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// GetSegments provides a mock function with given fields: id, segmentType, scope, collectionID +func (_m *ISegmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error) { + ret := _m.Called(id, segmentType, scope, collectionID) + + if len(ret) == 0 { + panic("no return value specified for GetSegments") + } + + var r0 []*dbmodel.SegmentAndMetadata + var r1 error + if rf, ok := ret.Get(0).(func(types.UniqueID, *string, *string, types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error)); ok { + return rf(id, segmentType, scope, collectionID) + } + if rf, ok := ret.Get(0).(func(types.UniqueID, *string, *string, types.UniqueID) []*dbmodel.SegmentAndMetadata); ok { + r0 = rf(id, segmentType, scope, collectionID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.SegmentAndMetadata) + } + } + + if rf, ok := ret.Get(1).(func(types.UniqueID, *string, *string, types.UniqueID) error); ok { + r1 = rf(id, segmentType, scope, collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: _a0 +func (_m *ISegmentDb) Insert(_a0 *dbmodel.Segment) error { + ret := _m.Called(_a0) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Segment) error); ok { + r0 = rf(_a0) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// RegisterFilePaths provides a mock function with given fields: flushSegmentCompactions +func (_m *ISegmentDb) RegisterFilePaths(flushSegmentCompactions []*model.FlushSegmentCompaction) error { + ret := _m.Called(flushSegmentCompactions) + + if len(ret) == 0 { + panic("no return value specified for RegisterFilePaths") + } + + var r0 error + if rf, ok := ret.Get(0).(func([]*model.FlushSegmentCompaction) error); ok { + r0 = rf(flushSegmentCompactions) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Update provides a mock function with given fields: _a0 +func (_m *ISegmentDb) Update(_a0 *dbmodel.UpdateSegment) error { + ret := _m.Called(_a0) + + if len(ret) == 0 { + panic("no return value specified for Update") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.UpdateSegment) error); ok { + r0 = rf(_a0) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewISegmentDb creates a new instance of ISegmentDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewISegmentDb(t interface { + mock.TestingT + Cleanup(func()) +}) *ISegmentDb { + mock := &ISegmentDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/ISegmentMetadataDb.go b/go/mocks/ISegmentMetadataDb.go new file mode 100644 index 00000000000..e65aa5ab3cc --- /dev/null +++ b/go/mocks/ISegmentMetadataDb.go @@ -0,0 +1,99 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// ISegmentMetadataDb is an autogenerated mock type for the ISegmentMetadataDb type +type ISegmentMetadataDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *ISegmentMetadataDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteBySegmentID provides a mock function with given fields: segmentID +func (_m *ISegmentMetadataDb) DeleteBySegmentID(segmentID string) error { + ret := _m.Called(segmentID) + + if len(ret) == 0 { + panic("no return value specified for DeleteBySegmentID") + } + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(segmentID) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteBySegmentIDAndKeys provides a mock function with given fields: segmentID, keys +func (_m *ISegmentMetadataDb) DeleteBySegmentIDAndKeys(segmentID string, keys []string) error { + ret := _m.Called(segmentID, keys) + + if len(ret) == 0 { + panic("no return value specified for DeleteBySegmentIDAndKeys") + } + + var r0 error + if rf, ok := ret.Get(0).(func(string, []string) error); ok { + r0 = rf(segmentID, keys) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Insert provides a mock function with given fields: in +func (_m *ISegmentMetadataDb) Insert(in []*dbmodel.SegmentMetadata) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func([]*dbmodel.SegmentMetadata) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewISegmentMetadataDb creates a new instance of ISegmentMetadataDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewISegmentMetadataDb(t interface { + mock.TestingT + Cleanup(func()) +}) *ISegmentMetadataDb { + mock := &ISegmentMetadataDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/ITenantDb.go b/go/mocks/ITenantDb.go new file mode 100644 index 00000000000..ffc9c9bb7df --- /dev/null +++ b/go/mocks/ITenantDb.go @@ -0,0 +1,171 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// ITenantDb is an autogenerated mock type for the ITenantDb type +type ITenantDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *ITenantDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// GetAllTenants provides a mock function with given fields: +func (_m *ITenantDb) GetAllTenants() ([]*dbmodel.Tenant, error) { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for GetAllTenants") + } + + var r0 []*dbmodel.Tenant + var r1 error + if rf, ok := ret.Get(0).(func() ([]*dbmodel.Tenant, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []*dbmodel.Tenant); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Tenant) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenants provides a mock function with given fields: tenantID +func (_m *ITenantDb) GetTenants(tenantID string) ([]*dbmodel.Tenant, error) { + ret := _m.Called(tenantID) + + if len(ret) == 0 { + panic("no return value specified for GetTenants") + } + + var r0 []*dbmodel.Tenant + var r1 error + if rf, ok := ret.Get(0).(func(string) ([]*dbmodel.Tenant, error)); ok { + return rf(tenantID) + } + if rf, ok := ret.Get(0).(func(string) []*dbmodel.Tenant); ok { + r0 = rf(tenantID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Tenant) + } + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(tenantID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenantsLastCompactionTime provides a mock function with given fields: tenantIDs +func (_m *ITenantDb) GetTenantsLastCompactionTime(tenantIDs []string) ([]*dbmodel.Tenant, error) { + ret := _m.Called(tenantIDs) + + if len(ret) == 0 { + panic("no return value specified for GetTenantsLastCompactionTime") + } + + var r0 []*dbmodel.Tenant + var r1 error + if rf, ok := ret.Get(0).(func([]string) ([]*dbmodel.Tenant, error)); ok { + return rf(tenantIDs) + } + if rf, ok := ret.Get(0).(func([]string) []*dbmodel.Tenant); ok { + r0 = rf(tenantIDs) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.Tenant) + } + } + + if rf, ok := ret.Get(1).(func([]string) error); ok { + r1 = rf(tenantIDs) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: in +func (_m *ITenantDb) Insert(in *dbmodel.Tenant) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Tenant) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// UpdateTenantLastCompactionTime provides a mock function with given fields: tenantID, lastCompactionTime +func (_m *ITenantDb) UpdateTenantLastCompactionTime(tenantID string, lastCompactionTime int64) error { + ret := _m.Called(tenantID, lastCompactionTime) + + if len(ret) == 0 { + panic("no return value specified for UpdateTenantLastCompactionTime") + } + + var r0 error + if rf, ok := ret.Get(0).(func(string, int64) error); ok { + r0 = rf(tenantID, lastCompactionTime) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewITenantDb creates a new instance of ITenantDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewITenantDb(t interface { + mock.TestingT + Cleanup(func()) +}) *ITenantDb { + mock := &ITenantDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/ITransaction.go b/go/mocks/ITransaction.go new file mode 100644 index 00000000000..884e3129bbd --- /dev/null +++ b/go/mocks/ITransaction.go @@ -0,0 +1,46 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + mock "github.com/stretchr/testify/mock" +) + +// ITransaction is an autogenerated mock type for the ITransaction type +type ITransaction struct { + mock.Mock +} + +// Transaction provides a mock function with given fields: ctx, fn +func (_m *ITransaction) Transaction(ctx context.Context, fn func(context.Context) error) error { + ret := _m.Called(ctx, fn) + + if len(ret) == 0 { + panic("no return value specified for Transaction") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, func(context.Context) error) error); ok { + r0 = rf(ctx, fn) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewITransaction creates a new instance of ITransaction. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewITransaction(t interface { + mock.TestingT + Cleanup(func()) +}) *ITransaction { + mock := &ITransaction{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/IWatcher.go b/go/mocks/IWatcher.go new file mode 100644 index 00000000000..eba7bd2520e --- /dev/null +++ b/go/mocks/IWatcher.go @@ -0,0 +1,96 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + memberlist_manager "github.com/chroma-core/chroma/go/pkg/memberlist_manager" + mock "github.com/stretchr/testify/mock" +) + +// IWatcher is an autogenerated mock type for the IWatcher type +type IWatcher struct { + mock.Mock +} + +// GetStatus provides a mock function with given fields: node_ip +func (_m *IWatcher) GetStatus(node_ip string) (memberlist_manager.Status, error) { + ret := _m.Called(node_ip) + + if len(ret) == 0 { + panic("no return value specified for GetStatus") + } + + var r0 memberlist_manager.Status + var r1 error + if rf, ok := ret.Get(0).(func(string) (memberlist_manager.Status, error)); ok { + return rf(node_ip) + } + if rf, ok := ret.Get(0).(func(string) memberlist_manager.Status); ok { + r0 = rf(node_ip) + } else { + r0 = ret.Get(0).(memberlist_manager.Status) + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(node_ip) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// RegisterCallback provides a mock function with given fields: callback +func (_m *IWatcher) RegisterCallback(callback memberlist_manager.NodeWatcherCallback) { + _m.Called(callback) +} + +// Start provides a mock function with given fields: +func (_m *IWatcher) Start() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Start") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Stop provides a mock function with given fields: +func (_m *IWatcher) Stop() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Stop") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewIWatcher creates a new instance of IWatcher. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewIWatcher(t interface { + mock.TestingT + Cleanup(func()) +}) *IWatcher { + mock := &IWatcher{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/LogServiceClient.go b/go/mocks/LogServiceClient.go new file mode 100644 index 00000000000..7f7a86cefa7 --- /dev/null +++ b/go/mocks/LogServiceClient.go @@ -0,0 +1,143 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + grpc "google.golang.org/grpc" + + logservicepb "github.com/chroma-core/chroma/go/pkg/proto/logservicepb" + + mock "github.com/stretchr/testify/mock" +) + +// LogServiceClient is an autogenerated mock type for the LogServiceClient type +type LogServiceClient struct { + mock.Mock +} + +// GetAllCollectionInfoToCompact provides a mock function with given fields: ctx, in, opts +func (_m *LogServiceClient) GetAllCollectionInfoToCompact(ctx context.Context, in *logservicepb.GetAllCollectionInfoToCompactRequest, opts ...grpc.CallOption) (*logservicepb.GetAllCollectionInfoToCompactResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetAllCollectionInfoToCompact") + } + + var r0 *logservicepb.GetAllCollectionInfoToCompactResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.GetAllCollectionInfoToCompactRequest, ...grpc.CallOption) (*logservicepb.GetAllCollectionInfoToCompactResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.GetAllCollectionInfoToCompactRequest, ...grpc.CallOption) *logservicepb.GetAllCollectionInfoToCompactResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*logservicepb.GetAllCollectionInfoToCompactResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *logservicepb.GetAllCollectionInfoToCompactRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PullLogs provides a mock function with given fields: ctx, in, opts +func (_m *LogServiceClient) PullLogs(ctx context.Context, in *logservicepb.PullLogsRequest, opts ...grpc.CallOption) (*logservicepb.PullLogsResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for PullLogs") + } + + var r0 *logservicepb.PullLogsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PullLogsRequest, ...grpc.CallOption) (*logservicepb.PullLogsResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PullLogsRequest, ...grpc.CallOption) *logservicepb.PullLogsResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*logservicepb.PullLogsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *logservicepb.PullLogsRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PushLogs provides a mock function with given fields: ctx, in, opts +func (_m *LogServiceClient) PushLogs(ctx context.Context, in *logservicepb.PushLogsRequest, opts ...grpc.CallOption) (*logservicepb.PushLogsResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for PushLogs") + } + + var r0 *logservicepb.PushLogsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PushLogsRequest, ...grpc.CallOption) (*logservicepb.PushLogsResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PushLogsRequest, ...grpc.CallOption) *logservicepb.PushLogsResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*logservicepb.PushLogsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *logservicepb.PushLogsRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewLogServiceClient creates a new instance of LogServiceClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewLogServiceClient(t interface { + mock.TestingT + Cleanup(func()) +}) *LogServiceClient { + mock := &LogServiceClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/LogServiceServer.go b/go/mocks/LogServiceServer.go new file mode 100644 index 00000000000..0cc7414f086 --- /dev/null +++ b/go/mocks/LogServiceServer.go @@ -0,0 +1,124 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + logservicepb "github.com/chroma-core/chroma/go/pkg/proto/logservicepb" + mock "github.com/stretchr/testify/mock" +) + +// LogServiceServer is an autogenerated mock type for the LogServiceServer type +type LogServiceServer struct { + mock.Mock +} + +// GetAllCollectionInfoToCompact provides a mock function with given fields: _a0, _a1 +func (_m *LogServiceServer) GetAllCollectionInfoToCompact(_a0 context.Context, _a1 *logservicepb.GetAllCollectionInfoToCompactRequest) (*logservicepb.GetAllCollectionInfoToCompactResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetAllCollectionInfoToCompact") + } + + var r0 *logservicepb.GetAllCollectionInfoToCompactResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.GetAllCollectionInfoToCompactRequest) (*logservicepb.GetAllCollectionInfoToCompactResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.GetAllCollectionInfoToCompactRequest) *logservicepb.GetAllCollectionInfoToCompactResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*logservicepb.GetAllCollectionInfoToCompactResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *logservicepb.GetAllCollectionInfoToCompactRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PullLogs provides a mock function with given fields: _a0, _a1 +func (_m *LogServiceServer) PullLogs(_a0 context.Context, _a1 *logservicepb.PullLogsRequest) (*logservicepb.PullLogsResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for PullLogs") + } + + var r0 *logservicepb.PullLogsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PullLogsRequest) (*logservicepb.PullLogsResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PullLogsRequest) *logservicepb.PullLogsResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*logservicepb.PullLogsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *logservicepb.PullLogsRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// PushLogs provides a mock function with given fields: _a0, _a1 +func (_m *LogServiceServer) PushLogs(_a0 context.Context, _a1 *logservicepb.PushLogsRequest) (*logservicepb.PushLogsResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for PushLogs") + } + + var r0 *logservicepb.PushLogsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PushLogsRequest) (*logservicepb.PushLogsResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *logservicepb.PushLogsRequest) *logservicepb.PushLogsResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*logservicepb.PushLogsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *logservicepb.PushLogsRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// mustEmbedUnimplementedLogServiceServer provides a mock function with given fields: +func (_m *LogServiceServer) mustEmbedUnimplementedLogServiceServer() { + _m.Called() +} + +// NewLogServiceServer creates a new instance of LogServiceServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewLogServiceServer(t interface { + mock.TestingT + Cleanup(func()) +}) *LogServiceServer { + mock := &LogServiceServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/NodeWatcherCallback.go b/go/mocks/NodeWatcherCallback.go new file mode 100644 index 00000000000..83af860bfc7 --- /dev/null +++ b/go/mocks/NodeWatcherCallback.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// NodeWatcherCallback is an autogenerated mock type for the NodeWatcherCallback type +type NodeWatcherCallback struct { + mock.Mock +} + +// Execute provides a mock function with given fields: node_ip +func (_m *NodeWatcherCallback) Execute(node_ip string) { + _m.Called(node_ip) +} + +// NewNodeWatcherCallback creates a new instance of NodeWatcherCallback. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewNodeWatcherCallback(t interface { + mock.TestingT + Cleanup(func()) +}) *NodeWatcherCallback { + mock := &NodeWatcherCallback{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/NotificationProcessor.go b/go/mocks/NotificationProcessor.go new file mode 100644 index 00000000000..381e5b748c8 --- /dev/null +++ b/go/mocks/NotificationProcessor.go @@ -0,0 +1,88 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + notification "github.com/chroma-core/chroma/go/pkg/notification" + mock "github.com/stretchr/testify/mock" +) + +// NotificationProcessor is an autogenerated mock type for the NotificationProcessor type +type NotificationProcessor struct { + mock.Mock +} + +// Process provides a mock function with given fields: ctx +func (_m *NotificationProcessor) Process(ctx context.Context) error { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for Process") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context) error); ok { + r0 = rf(ctx) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Start provides a mock function with given fields: +func (_m *NotificationProcessor) Start() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Start") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Stop provides a mock function with given fields: +func (_m *NotificationProcessor) Stop() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Stop") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Trigger provides a mock function with given fields: ctx, triggerMsg +func (_m *NotificationProcessor) Trigger(ctx context.Context, triggerMsg notification.TriggerMessage) { + _m.Called(ctx, triggerMsg) +} + +// NewNotificationProcessor creates a new instance of NotificationProcessor. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewNotificationProcessor(t interface { + mock.TestingT + Cleanup(func()) +}) *NotificationProcessor { + mock := &NotificationProcessor{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/NotificationStore.go b/go/mocks/NotificationStore.go new file mode 100644 index 00000000000..9a744edc8a6 --- /dev/null +++ b/go/mocks/NotificationStore.go @@ -0,0 +1,125 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + model "github.com/chroma-core/chroma/go/pkg/model" + mock "github.com/stretchr/testify/mock" +) + +// NotificationStore is an autogenerated mock type for the NotificationStore type +type NotificationStore struct { + mock.Mock +} + +// AddNotification provides a mock function with given fields: ctx, _a1 +func (_m *NotificationStore) AddNotification(ctx context.Context, _a1 model.Notification) error { + ret := _m.Called(ctx, _a1) + + if len(ret) == 0 { + panic("no return value specified for AddNotification") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, model.Notification) error); ok { + r0 = rf(ctx, _a1) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// GetAllPendingNotifications provides a mock function with given fields: ctx +func (_m *NotificationStore) GetAllPendingNotifications(ctx context.Context) (map[string][]model.Notification, error) { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for GetAllPendingNotifications") + } + + var r0 map[string][]model.Notification + var r1 error + if rf, ok := ret.Get(0).(func(context.Context) (map[string][]model.Notification, error)); ok { + return rf(ctx) + } + if rf, ok := ret.Get(0).(func(context.Context) map[string][]model.Notification); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(map[string][]model.Notification) + } + } + + if rf, ok := ret.Get(1).(func(context.Context) error); ok { + r1 = rf(ctx) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetNotifications provides a mock function with given fields: ctx, collecitonID +func (_m *NotificationStore) GetNotifications(ctx context.Context, collecitonID string) ([]model.Notification, error) { + ret := _m.Called(ctx, collecitonID) + + if len(ret) == 0 { + panic("no return value specified for GetNotifications") + } + + var r0 []model.Notification + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string) ([]model.Notification, error)); ok { + return rf(ctx, collecitonID) + } + if rf, ok := ret.Get(0).(func(context.Context, string) []model.Notification); ok { + r0 = rf(ctx, collecitonID) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]model.Notification) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, collecitonID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// RemoveNotifications provides a mock function with given fields: ctx, notifications +func (_m *NotificationStore) RemoveNotifications(ctx context.Context, notifications []model.Notification) error { + ret := _m.Called(ctx, notifications) + + if len(ret) == 0 { + panic("no return value specified for RemoveNotifications") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, []model.Notification) error); ok { + r0 = rf(ctx, notifications) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewNotificationStore creates a new instance of NotificationStore. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewNotificationStore(t interface { + mock.TestingT + Cleanup(func()) +}) *NotificationStore { + mock := &NotificationStore{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/Notifier.go b/go/mocks/Notifier.go new file mode 100644 index 00000000000..2d2be648291 --- /dev/null +++ b/go/mocks/Notifier.go @@ -0,0 +1,47 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + model "github.com/chroma-core/chroma/go/pkg/model" + mock "github.com/stretchr/testify/mock" +) + +// Notifier is an autogenerated mock type for the Notifier type +type Notifier struct { + mock.Mock +} + +// Notify provides a mock function with given fields: ctx, notifications +func (_m *Notifier) Notify(ctx context.Context, notifications []model.Notification) error { + ret := _m.Called(ctx, notifications) + + if len(ret) == 0 { + panic("no return value specified for Notify") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, []model.Notification) error); ok { + r0 = rf(ctx, notifications) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewNotifier creates a new instance of Notifier. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewNotifier(t interface { + mock.TestingT + Cleanup(func()) +}) *Notifier { + mock := &Notifier{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/SegmentMetadataValueType.go b/go/mocks/SegmentMetadataValueType.go new file mode 100644 index 00000000000..f742156d07d --- /dev/null +++ b/go/mocks/SegmentMetadataValueType.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// SegmentMetadataValueType is an autogenerated mock type for the SegmentMetadataValueType type +type SegmentMetadataValueType struct { + mock.Mock +} + +// IsSegmentMetadataValueType provides a mock function with given fields: +func (_m *SegmentMetadataValueType) IsSegmentMetadataValueType() { + _m.Called() +} + +// NewSegmentMetadataValueType creates a new instance of SegmentMetadataValueType. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewSegmentMetadataValueType(t interface { + mock.TestingT + Cleanup(func()) +}) *SegmentMetadataValueType { + mock := &SegmentMetadataValueType{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/SysDBClient.go b/go/mocks/SysDBClient.go new file mode 100644 index 00000000000..c2eeb3bc44e --- /dev/null +++ b/go/mocks/SysDBClient.go @@ -0,0 +1,625 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + coordinatorpb "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + emptypb "google.golang.org/protobuf/types/known/emptypb" + + grpc "google.golang.org/grpc" + + mock "github.com/stretchr/testify/mock" +) + +// SysDBClient is an autogenerated mock type for the SysDBClient type +type SysDBClient struct { + mock.Mock +} + +// CreateCollection provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) CreateCollection(ctx context.Context, in *coordinatorpb.CreateCollectionRequest, opts ...grpc.CallOption) (*coordinatorpb.CreateCollectionResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for CreateCollection") + } + + var r0 *coordinatorpb.CreateCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateCollectionRequest, ...grpc.CallOption) (*coordinatorpb.CreateCollectionResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateCollectionRequest, ...grpc.CallOption) *coordinatorpb.CreateCollectionResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateCollectionRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateDatabase provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) CreateDatabase(ctx context.Context, in *coordinatorpb.CreateDatabaseRequest, opts ...grpc.CallOption) (*coordinatorpb.CreateDatabaseResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for CreateDatabase") + } + + var r0 *coordinatorpb.CreateDatabaseResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateDatabaseRequest, ...grpc.CallOption) (*coordinatorpb.CreateDatabaseResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateDatabaseRequest, ...grpc.CallOption) *coordinatorpb.CreateDatabaseResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateDatabaseResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateDatabaseRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateSegment provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) CreateSegment(ctx context.Context, in *coordinatorpb.CreateSegmentRequest, opts ...grpc.CallOption) (*coordinatorpb.CreateSegmentResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for CreateSegment") + } + + var r0 *coordinatorpb.CreateSegmentResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateSegmentRequest, ...grpc.CallOption) (*coordinatorpb.CreateSegmentResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateSegmentRequest, ...grpc.CallOption) *coordinatorpb.CreateSegmentResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateSegmentResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateSegmentRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateTenant provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) CreateTenant(ctx context.Context, in *coordinatorpb.CreateTenantRequest, opts ...grpc.CallOption) (*coordinatorpb.CreateTenantResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for CreateTenant") + } + + var r0 *coordinatorpb.CreateTenantResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateTenantRequest, ...grpc.CallOption) (*coordinatorpb.CreateTenantResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateTenantRequest, ...grpc.CallOption) *coordinatorpb.CreateTenantResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateTenantResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateTenantRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// DeleteCollection provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) DeleteCollection(ctx context.Context, in *coordinatorpb.DeleteCollectionRequest, opts ...grpc.CallOption) (*coordinatorpb.DeleteCollectionResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for DeleteCollection") + } + + var r0 *coordinatorpb.DeleteCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteCollectionRequest, ...grpc.CallOption) (*coordinatorpb.DeleteCollectionResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteCollectionRequest, ...grpc.CallOption) *coordinatorpb.DeleteCollectionResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.DeleteCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.DeleteCollectionRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// DeleteSegment provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) DeleteSegment(ctx context.Context, in *coordinatorpb.DeleteSegmentRequest, opts ...grpc.CallOption) (*coordinatorpb.DeleteSegmentResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for DeleteSegment") + } + + var r0 *coordinatorpb.DeleteSegmentResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteSegmentRequest, ...grpc.CallOption) (*coordinatorpb.DeleteSegmentResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteSegmentRequest, ...grpc.CallOption) *coordinatorpb.DeleteSegmentResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.DeleteSegmentResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.DeleteSegmentRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// FlushCollectionCompaction provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) FlushCollectionCompaction(ctx context.Context, in *coordinatorpb.FlushCollectionCompactionRequest, opts ...grpc.CallOption) (*coordinatorpb.FlushCollectionCompactionResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for FlushCollectionCompaction") + } + + var r0 *coordinatorpb.FlushCollectionCompactionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.FlushCollectionCompactionRequest, ...grpc.CallOption) (*coordinatorpb.FlushCollectionCompactionResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.FlushCollectionCompactionRequest, ...grpc.CallOption) *coordinatorpb.FlushCollectionCompactionResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.FlushCollectionCompactionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.FlushCollectionCompactionRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetCollections provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) GetCollections(ctx context.Context, in *coordinatorpb.GetCollectionsRequest, opts ...grpc.CallOption) (*coordinatorpb.GetCollectionsResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetCollections") + } + + var r0 *coordinatorpb.GetCollectionsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetCollectionsRequest, ...grpc.CallOption) (*coordinatorpb.GetCollectionsResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetCollectionsRequest, ...grpc.CallOption) *coordinatorpb.GetCollectionsResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetCollectionsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetCollectionsRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetDatabase provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) GetDatabase(ctx context.Context, in *coordinatorpb.GetDatabaseRequest, opts ...grpc.CallOption) (*coordinatorpb.GetDatabaseResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetDatabase") + } + + var r0 *coordinatorpb.GetDatabaseResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetDatabaseRequest, ...grpc.CallOption) (*coordinatorpb.GetDatabaseResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetDatabaseRequest, ...grpc.CallOption) *coordinatorpb.GetDatabaseResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetDatabaseResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetDatabaseRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetLastCompactionTimeForTenant provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) GetLastCompactionTimeForTenant(ctx context.Context, in *coordinatorpb.GetLastCompactionTimeForTenantRequest, opts ...grpc.CallOption) (*coordinatorpb.GetLastCompactionTimeForTenantResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetLastCompactionTimeForTenant") + } + + var r0 *coordinatorpb.GetLastCompactionTimeForTenantResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetLastCompactionTimeForTenantRequest, ...grpc.CallOption) (*coordinatorpb.GetLastCompactionTimeForTenantResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetLastCompactionTimeForTenantRequest, ...grpc.CallOption) *coordinatorpb.GetLastCompactionTimeForTenantResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetLastCompactionTimeForTenantResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetLastCompactionTimeForTenantRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetSegments provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) GetSegments(ctx context.Context, in *coordinatorpb.GetSegmentsRequest, opts ...grpc.CallOption) (*coordinatorpb.GetSegmentsResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetSegments") + } + + var r0 *coordinatorpb.GetSegmentsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetSegmentsRequest, ...grpc.CallOption) (*coordinatorpb.GetSegmentsResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetSegmentsRequest, ...grpc.CallOption) *coordinatorpb.GetSegmentsResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetSegmentsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetSegmentsRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenant provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) GetTenant(ctx context.Context, in *coordinatorpb.GetTenantRequest, opts ...grpc.CallOption) (*coordinatorpb.GetTenantResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetTenant") + } + + var r0 *coordinatorpb.GetTenantResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetTenantRequest, ...grpc.CallOption) (*coordinatorpb.GetTenantResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetTenantRequest, ...grpc.CallOption) *coordinatorpb.GetTenantResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetTenantResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetTenantRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// ResetState provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) ResetState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*coordinatorpb.ResetStateResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for ResetState") + } + + var r0 *coordinatorpb.ResetStateResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *emptypb.Empty, ...grpc.CallOption) (*coordinatorpb.ResetStateResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *emptypb.Empty, ...grpc.CallOption) *coordinatorpb.ResetStateResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.ResetStateResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *emptypb.Empty, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// SetLastCompactionTimeForTenant provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) SetLastCompactionTimeForTenant(ctx context.Context, in *coordinatorpb.SetLastCompactionTimeForTenantRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for SetLastCompactionTimeForTenant") + } + + var r0 *emptypb.Empty + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.SetLastCompactionTimeForTenantRequest, ...grpc.CallOption) (*emptypb.Empty, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.SetLastCompactionTimeForTenantRequest, ...grpc.CallOption) *emptypb.Empty); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*emptypb.Empty) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.SetLastCompactionTimeForTenantRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// UpdateCollection provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) UpdateCollection(ctx context.Context, in *coordinatorpb.UpdateCollectionRequest, opts ...grpc.CallOption) (*coordinatorpb.UpdateCollectionResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for UpdateCollection") + } + + var r0 *coordinatorpb.UpdateCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateCollectionRequest, ...grpc.CallOption) (*coordinatorpb.UpdateCollectionResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateCollectionRequest, ...grpc.CallOption) *coordinatorpb.UpdateCollectionResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.UpdateCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.UpdateCollectionRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// UpdateSegment provides a mock function with given fields: ctx, in, opts +func (_m *SysDBClient) UpdateSegment(ctx context.Context, in *coordinatorpb.UpdateSegmentRequest, opts ...grpc.CallOption) (*coordinatorpb.UpdateSegmentResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for UpdateSegment") + } + + var r0 *coordinatorpb.UpdateSegmentResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateSegmentRequest, ...grpc.CallOption) (*coordinatorpb.UpdateSegmentResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateSegmentRequest, ...grpc.CallOption) *coordinatorpb.UpdateSegmentResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.UpdateSegmentResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.UpdateSegmentRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewSysDBClient creates a new instance of SysDBClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewSysDBClient(t interface { + mock.TestingT + Cleanup(func()) +}) *SysDBClient { + mock := &SysDBClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/SysDBServer.go b/go/mocks/SysDBServer.go new file mode 100644 index 00000000000..45b53925cdf --- /dev/null +++ b/go/mocks/SysDBServer.go @@ -0,0 +1,516 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + coordinatorpb "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + emptypb "google.golang.org/protobuf/types/known/emptypb" + + mock "github.com/stretchr/testify/mock" +) + +// SysDBServer is an autogenerated mock type for the SysDBServer type +type SysDBServer struct { + mock.Mock +} + +// CreateCollection provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) CreateCollection(_a0 context.Context, _a1 *coordinatorpb.CreateCollectionRequest) (*coordinatorpb.CreateCollectionResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for CreateCollection") + } + + var r0 *coordinatorpb.CreateCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateCollectionRequest) (*coordinatorpb.CreateCollectionResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateCollectionRequest) *coordinatorpb.CreateCollectionResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateCollectionRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateDatabase provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) CreateDatabase(_a0 context.Context, _a1 *coordinatorpb.CreateDatabaseRequest) (*coordinatorpb.CreateDatabaseResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for CreateDatabase") + } + + var r0 *coordinatorpb.CreateDatabaseResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateDatabaseRequest) (*coordinatorpb.CreateDatabaseResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateDatabaseRequest) *coordinatorpb.CreateDatabaseResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateDatabaseResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateDatabaseRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateSegment provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) CreateSegment(_a0 context.Context, _a1 *coordinatorpb.CreateSegmentRequest) (*coordinatorpb.CreateSegmentResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for CreateSegment") + } + + var r0 *coordinatorpb.CreateSegmentResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateSegmentRequest) (*coordinatorpb.CreateSegmentResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateSegmentRequest) *coordinatorpb.CreateSegmentResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateSegmentResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateSegmentRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CreateTenant provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) CreateTenant(_a0 context.Context, _a1 *coordinatorpb.CreateTenantRequest) (*coordinatorpb.CreateTenantResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for CreateTenant") + } + + var r0 *coordinatorpb.CreateTenantResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateTenantRequest) (*coordinatorpb.CreateTenantResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.CreateTenantRequest) *coordinatorpb.CreateTenantResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.CreateTenantResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.CreateTenantRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// DeleteCollection provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) DeleteCollection(_a0 context.Context, _a1 *coordinatorpb.DeleteCollectionRequest) (*coordinatorpb.DeleteCollectionResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for DeleteCollection") + } + + var r0 *coordinatorpb.DeleteCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteCollectionRequest) (*coordinatorpb.DeleteCollectionResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteCollectionRequest) *coordinatorpb.DeleteCollectionResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.DeleteCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.DeleteCollectionRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// DeleteSegment provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) DeleteSegment(_a0 context.Context, _a1 *coordinatorpb.DeleteSegmentRequest) (*coordinatorpb.DeleteSegmentResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for DeleteSegment") + } + + var r0 *coordinatorpb.DeleteSegmentResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteSegmentRequest) (*coordinatorpb.DeleteSegmentResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.DeleteSegmentRequest) *coordinatorpb.DeleteSegmentResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.DeleteSegmentResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.DeleteSegmentRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// FlushCollectionCompaction provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) FlushCollectionCompaction(_a0 context.Context, _a1 *coordinatorpb.FlushCollectionCompactionRequest) (*coordinatorpb.FlushCollectionCompactionResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for FlushCollectionCompaction") + } + + var r0 *coordinatorpb.FlushCollectionCompactionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.FlushCollectionCompactionRequest) (*coordinatorpb.FlushCollectionCompactionResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.FlushCollectionCompactionRequest) *coordinatorpb.FlushCollectionCompactionResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.FlushCollectionCompactionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.FlushCollectionCompactionRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetCollections provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) GetCollections(_a0 context.Context, _a1 *coordinatorpb.GetCollectionsRequest) (*coordinatorpb.GetCollectionsResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetCollections") + } + + var r0 *coordinatorpb.GetCollectionsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetCollectionsRequest) (*coordinatorpb.GetCollectionsResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetCollectionsRequest) *coordinatorpb.GetCollectionsResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetCollectionsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetCollectionsRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetDatabase provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) GetDatabase(_a0 context.Context, _a1 *coordinatorpb.GetDatabaseRequest) (*coordinatorpb.GetDatabaseResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetDatabase") + } + + var r0 *coordinatorpb.GetDatabaseResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetDatabaseRequest) (*coordinatorpb.GetDatabaseResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetDatabaseRequest) *coordinatorpb.GetDatabaseResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetDatabaseResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetDatabaseRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetLastCompactionTimeForTenant provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) GetLastCompactionTimeForTenant(_a0 context.Context, _a1 *coordinatorpb.GetLastCompactionTimeForTenantRequest) (*coordinatorpb.GetLastCompactionTimeForTenantResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetLastCompactionTimeForTenant") + } + + var r0 *coordinatorpb.GetLastCompactionTimeForTenantResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetLastCompactionTimeForTenantRequest) (*coordinatorpb.GetLastCompactionTimeForTenantResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetLastCompactionTimeForTenantRequest) *coordinatorpb.GetLastCompactionTimeForTenantResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetLastCompactionTimeForTenantResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetLastCompactionTimeForTenantRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetSegments provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) GetSegments(_a0 context.Context, _a1 *coordinatorpb.GetSegmentsRequest) (*coordinatorpb.GetSegmentsResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetSegments") + } + + var r0 *coordinatorpb.GetSegmentsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetSegmentsRequest) (*coordinatorpb.GetSegmentsResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetSegmentsRequest) *coordinatorpb.GetSegmentsResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetSegmentsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetSegmentsRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetTenant provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) GetTenant(_a0 context.Context, _a1 *coordinatorpb.GetTenantRequest) (*coordinatorpb.GetTenantResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetTenant") + } + + var r0 *coordinatorpb.GetTenantResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetTenantRequest) (*coordinatorpb.GetTenantResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetTenantRequest) *coordinatorpb.GetTenantResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetTenantResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetTenantRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// ResetState provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) ResetState(_a0 context.Context, _a1 *emptypb.Empty) (*coordinatorpb.ResetStateResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for ResetState") + } + + var r0 *coordinatorpb.ResetStateResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *emptypb.Empty) (*coordinatorpb.ResetStateResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *emptypb.Empty) *coordinatorpb.ResetStateResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.ResetStateResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *emptypb.Empty) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// SetLastCompactionTimeForTenant provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) SetLastCompactionTimeForTenant(_a0 context.Context, _a1 *coordinatorpb.SetLastCompactionTimeForTenantRequest) (*emptypb.Empty, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for SetLastCompactionTimeForTenant") + } + + var r0 *emptypb.Empty + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.SetLastCompactionTimeForTenantRequest) (*emptypb.Empty, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.SetLastCompactionTimeForTenantRequest) *emptypb.Empty); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*emptypb.Empty) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.SetLastCompactionTimeForTenantRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// UpdateCollection provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) UpdateCollection(_a0 context.Context, _a1 *coordinatorpb.UpdateCollectionRequest) (*coordinatorpb.UpdateCollectionResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for UpdateCollection") + } + + var r0 *coordinatorpb.UpdateCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateCollectionRequest) (*coordinatorpb.UpdateCollectionResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateCollectionRequest) *coordinatorpb.UpdateCollectionResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.UpdateCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.UpdateCollectionRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// UpdateSegment provides a mock function with given fields: _a0, _a1 +func (_m *SysDBServer) UpdateSegment(_a0 context.Context, _a1 *coordinatorpb.UpdateSegmentRequest) (*coordinatorpb.UpdateSegmentResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for UpdateSegment") + } + + var r0 *coordinatorpb.UpdateSegmentResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateSegmentRequest) (*coordinatorpb.UpdateSegmentResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.UpdateSegmentRequest) *coordinatorpb.UpdateSegmentResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.UpdateSegmentResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.UpdateSegmentRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// mustEmbedUnimplementedSysDBServer provides a mock function with given fields: +func (_m *SysDBServer) mustEmbedUnimplementedSysDBServer() { + _m.Called() +} + +// NewSysDBServer creates a new instance of SysDBServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewSysDBServer(t interface { + mock.TestingT + Cleanup(func()) +}) *SysDBServer { + mock := &SysDBServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/UnsafeLogServiceServer.go b/go/mocks/UnsafeLogServiceServer.go new file mode 100644 index 00000000000..92a15424ae0 --- /dev/null +++ b/go/mocks/UnsafeLogServiceServer.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// UnsafeLogServiceServer is an autogenerated mock type for the UnsafeLogServiceServer type +type UnsafeLogServiceServer struct { + mock.Mock +} + +// mustEmbedUnimplementedLogServiceServer provides a mock function with given fields: +func (_m *UnsafeLogServiceServer) mustEmbedUnimplementedLogServiceServer() { + _m.Called() +} + +// NewUnsafeLogServiceServer creates a new instance of UnsafeLogServiceServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewUnsafeLogServiceServer(t interface { + mock.TestingT + Cleanup(func()) +}) *UnsafeLogServiceServer { + mock := &UnsafeLogServiceServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/UnsafeSysDBServer.go b/go/mocks/UnsafeSysDBServer.go new file mode 100644 index 00000000000..a45b6af2001 --- /dev/null +++ b/go/mocks/UnsafeSysDBServer.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// UnsafeSysDBServer is an autogenerated mock type for the UnsafeSysDBServer type +type UnsafeSysDBServer struct { + mock.Mock +} + +// mustEmbedUnimplementedSysDBServer provides a mock function with given fields: +func (_m *UnsafeSysDBServer) mustEmbedUnimplementedSysDBServer() { + _m.Called() +} + +// NewUnsafeSysDBServer creates a new instance of UnsafeSysDBServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewUnsafeSysDBServer(t interface { + mock.TestingT + Cleanup(func()) +}) *UnsafeSysDBServer { + mock := &UnsafeSysDBServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/UnsafeVectorReaderServer.go b/go/mocks/UnsafeVectorReaderServer.go new file mode 100644 index 00000000000..a55c0a1663c --- /dev/null +++ b/go/mocks/UnsafeVectorReaderServer.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// UnsafeVectorReaderServer is an autogenerated mock type for the UnsafeVectorReaderServer type +type UnsafeVectorReaderServer struct { + mock.Mock +} + +// mustEmbedUnimplementedVectorReaderServer provides a mock function with given fields: +func (_m *UnsafeVectorReaderServer) mustEmbedUnimplementedVectorReaderServer() { + _m.Called() +} + +// NewUnsafeVectorReaderServer creates a new instance of UnsafeVectorReaderServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewUnsafeVectorReaderServer(t interface { + mock.TestingT + Cleanup(func()) +}) *UnsafeVectorReaderServer { + mock := &UnsafeVectorReaderServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/VectorReaderClient.go b/go/mocks/VectorReaderClient.go new file mode 100644 index 00000000000..62c436f4670 --- /dev/null +++ b/go/mocks/VectorReaderClient.go @@ -0,0 +1,105 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + coordinatorpb "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + grpc "google.golang.org/grpc" + + mock "github.com/stretchr/testify/mock" +) + +// VectorReaderClient is an autogenerated mock type for the VectorReaderClient type +type VectorReaderClient struct { + mock.Mock +} + +// GetVectors provides a mock function with given fields: ctx, in, opts +func (_m *VectorReaderClient) GetVectors(ctx context.Context, in *coordinatorpb.GetVectorsRequest, opts ...grpc.CallOption) (*coordinatorpb.GetVectorsResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for GetVectors") + } + + var r0 *coordinatorpb.GetVectorsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetVectorsRequest, ...grpc.CallOption) (*coordinatorpb.GetVectorsResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetVectorsRequest, ...grpc.CallOption) *coordinatorpb.GetVectorsResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetVectorsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetVectorsRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// QueryVectors provides a mock function with given fields: ctx, in, opts +func (_m *VectorReaderClient) QueryVectors(ctx context.Context, in *coordinatorpb.QueryVectorsRequest, opts ...grpc.CallOption) (*coordinatorpb.QueryVectorsResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for QueryVectors") + } + + var r0 *coordinatorpb.QueryVectorsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.QueryVectorsRequest, ...grpc.CallOption) (*coordinatorpb.QueryVectorsResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.QueryVectorsRequest, ...grpc.CallOption) *coordinatorpb.QueryVectorsResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.QueryVectorsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.QueryVectorsRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewVectorReaderClient creates a new instance of VectorReaderClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewVectorReaderClient(t interface { + mock.TestingT + Cleanup(func()) +}) *VectorReaderClient { + mock := &VectorReaderClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/VectorReaderServer.go b/go/mocks/VectorReaderServer.go new file mode 100644 index 00000000000..55852d2e43b --- /dev/null +++ b/go/mocks/VectorReaderServer.go @@ -0,0 +1,94 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import ( + context "context" + + coordinatorpb "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + mock "github.com/stretchr/testify/mock" +) + +// VectorReaderServer is an autogenerated mock type for the VectorReaderServer type +type VectorReaderServer struct { + mock.Mock +} + +// GetVectors provides a mock function with given fields: _a0, _a1 +func (_m *VectorReaderServer) GetVectors(_a0 context.Context, _a1 *coordinatorpb.GetVectorsRequest) (*coordinatorpb.GetVectorsResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for GetVectors") + } + + var r0 *coordinatorpb.GetVectorsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetVectorsRequest) (*coordinatorpb.GetVectorsResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.GetVectorsRequest) *coordinatorpb.GetVectorsResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.GetVectorsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.GetVectorsRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// QueryVectors provides a mock function with given fields: _a0, _a1 +func (_m *VectorReaderServer) QueryVectors(_a0 context.Context, _a1 *coordinatorpb.QueryVectorsRequest) (*coordinatorpb.QueryVectorsResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for QueryVectors") + } + + var r0 *coordinatorpb.QueryVectorsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.QueryVectorsRequest) (*coordinatorpb.QueryVectorsResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *coordinatorpb.QueryVectorsRequest) *coordinatorpb.QueryVectorsResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*coordinatorpb.QueryVectorsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *coordinatorpb.QueryVectorsRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// mustEmbedUnimplementedVectorReaderServer provides a mock function with given fields: +func (_m *VectorReaderServer) mustEmbedUnimplementedVectorReaderServer() { + _m.Called() +} + +// NewVectorReaderServer creates a new instance of VectorReaderServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewVectorReaderServer(t interface { + mock.TestingT + Cleanup(func()) +}) *VectorReaderServer { + mock := &VectorReaderServer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/isUpdateCollectionRequest_MetadataUpdate.go b/go/mocks/isUpdateCollectionRequest_MetadataUpdate.go new file mode 100644 index 00000000000..204cb9a51a8 --- /dev/null +++ b/go/mocks/isUpdateCollectionRequest_MetadataUpdate.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// isUpdateCollectionRequest_MetadataUpdate is an autogenerated mock type for the isUpdateCollectionRequest_MetadataUpdate type +type isUpdateCollectionRequest_MetadataUpdate struct { + mock.Mock +} + +// isUpdateCollectionRequest_MetadataUpdate provides a mock function with given fields: +func (_m *isUpdateCollectionRequest_MetadataUpdate) isUpdateCollectionRequest_MetadataUpdate() { + _m.Called() +} + +// newIsUpdateCollectionRequest_MetadataUpdate creates a new instance of isUpdateCollectionRequest_MetadataUpdate. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func newIsUpdateCollectionRequest_MetadataUpdate(t interface { + mock.TestingT + Cleanup(func()) +}) *isUpdateCollectionRequest_MetadataUpdate { + mock := &isUpdateCollectionRequest_MetadataUpdate{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/isUpdateMetadataValue_Value.go b/go/mocks/isUpdateMetadataValue_Value.go new file mode 100644 index 00000000000..51187097971 --- /dev/null +++ b/go/mocks/isUpdateMetadataValue_Value.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// isUpdateMetadataValue_Value is an autogenerated mock type for the isUpdateMetadataValue_Value type +type isUpdateMetadataValue_Value struct { + mock.Mock +} + +// isUpdateMetadataValue_Value provides a mock function with given fields: +func (_m *isUpdateMetadataValue_Value) isUpdateMetadataValue_Value() { + _m.Called() +} + +// newIsUpdateMetadataValue_Value creates a new instance of isUpdateMetadataValue_Value. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func newIsUpdateMetadataValue_Value(t interface { + mock.TestingT + Cleanup(func()) +}) *isUpdateMetadataValue_Value { + mock := &isUpdateMetadataValue_Value{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/isUpdateSegmentRequest_CollectionUpdate.go b/go/mocks/isUpdateSegmentRequest_CollectionUpdate.go new file mode 100644 index 00000000000..9a1dee41326 --- /dev/null +++ b/go/mocks/isUpdateSegmentRequest_CollectionUpdate.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// isUpdateSegmentRequest_CollectionUpdate is an autogenerated mock type for the isUpdateSegmentRequest_CollectionUpdate type +type isUpdateSegmentRequest_CollectionUpdate struct { + mock.Mock +} + +// isUpdateSegmentRequest_CollectionUpdate provides a mock function with given fields: +func (_m *isUpdateSegmentRequest_CollectionUpdate) isUpdateSegmentRequest_CollectionUpdate() { + _m.Called() +} + +// newIsUpdateSegmentRequest_CollectionUpdate creates a new instance of isUpdateSegmentRequest_CollectionUpdate. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func newIsUpdateSegmentRequest_CollectionUpdate(t interface { + mock.TestingT + Cleanup(func()) +}) *isUpdateSegmentRequest_CollectionUpdate { + mock := &isUpdateSegmentRequest_CollectionUpdate{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/mocks/isUpdateSegmentRequest_MetadataUpdate.go b/go/mocks/isUpdateSegmentRequest_MetadataUpdate.go new file mode 100644 index 00000000000..15fffa75bee --- /dev/null +++ b/go/mocks/isUpdateSegmentRequest_MetadataUpdate.go @@ -0,0 +1,29 @@ +// Code generated by mockery v2.42.1. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// isUpdateSegmentRequest_MetadataUpdate is an autogenerated mock type for the isUpdateSegmentRequest_MetadataUpdate type +type isUpdateSegmentRequest_MetadataUpdate struct { + mock.Mock +} + +// isUpdateSegmentRequest_MetadataUpdate provides a mock function with given fields: +func (_m *isUpdateSegmentRequest_MetadataUpdate) isUpdateSegmentRequest_MetadataUpdate() { + _m.Called() +} + +// newIsUpdateSegmentRequest_MetadataUpdate creates a new instance of isUpdateSegmentRequest_MetadataUpdate. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func newIsUpdateSegmentRequest_MetadataUpdate(t interface { + mock.TestingT + Cleanup(func()) +}) *isUpdateSegmentRequest_MetadataUpdate { + mock := &isUpdateSegmentRequest_MetadataUpdate{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/coordinator/internal/common/component.go b/go/pkg/common/component.go similarity index 100% rename from go/coordinator/internal/common/component.go rename to go/pkg/common/component.go diff --git a/go/coordinator/internal/common/constants.go b/go/pkg/common/constants.go similarity index 100% rename from go/coordinator/internal/common/constants.go rename to go/pkg/common/constants.go diff --git a/go/coordinator/internal/common/errors.go b/go/pkg/common/errors.go similarity index 82% rename from go/coordinator/internal/common/errors.go rename to go/pkg/common/errors.go index 0275e2b6574..53da456208f 100644 --- a/go/coordinator/internal/common/errors.go +++ b/go/pkg/common/errors.go @@ -17,9 +17,11 @@ var ( ErrCollectionNotFound = errors.New("collection not found") ErrCollectionIDFormat = errors.New("collection id format error") ErrCollectionNameEmpty = errors.New("collection name is empty") - ErrCollectionTopicEmpty = errors.New("collection topic is empty") ErrCollectionUniqueConstraintViolation = errors.New("collection unique constraint violation") ErrCollectionDeleteNonExistingCollection = errors.New("delete non existing collection") + ErrCollectionLogPositionStale = errors.New("collection log position Stale") + ErrCollectionVersionStale = errors.New("collection version stale") + ErrCollectionVersionInvalid = errors.New("collection version invalid") // Collection metadata errors ErrUnknownCollectionMetadataType = errors.New("collection metadata value type not supported") @@ -27,10 +29,10 @@ var ( // Segment errors ErrSegmentIDFormat = errors.New("segment id format error") - ErrInvalidTopicUpdate = errors.New("invalid topic update, reset topic true and topic value not empty") ErrInvalidCollectionUpdate = errors.New("invalid collection update, reset collection true and collection value not empty") ErrSegmentUniqueConstraintViolation = errors.New("unique constraint violation") ErrSegmentDeleteNonExistingSegment = errors.New("delete non existing segment") + ErrSegmentUpdateNonExistingSegment = errors.New("update non existing segment") // Segment metadata errors ErrUnknownSegmentMetadataType = errors.New("segment metadata value type not supported") diff --git a/go/coordinator/internal/coordinator/apis.go b/go/pkg/coordinator/apis.go similarity index 62% rename from go/coordinator/internal/coordinator/apis.go rename to go/pkg/coordinator/apis.go index 24cb1a5ee13..a45dab0102c 100644 --- a/go/coordinator/internal/coordinator/apis.go +++ b/go/pkg/coordinator/apis.go @@ -2,11 +2,11 @@ package coordinator import ( "context" - "errors" - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" ) @@ -18,25 +18,28 @@ type ICoordinator interface { common.Component ResetState(ctx context.Context) error CreateCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error) - GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, dataName string) ([]*model.Collection, error) + GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, tenantID string, dataName string) ([]*model.Collection, error) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection) (*model.Collection, error) CreateSegment(ctx context.Context, createSegment *model.CreateSegment) error - GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*model.Segment, error) + GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*model.Segment, error) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment) (*model.Segment, error) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error) GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error) CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error) + SetTenantLastCompactionTime(ctx context.Context, tenantID string, lastCompactionTime int64) error + GetTenantsLastCompactionTime(ctx context.Context, tenantIDs []string) ([]*dbmodel.Tenant, error) + FlushCollectionCompaction(ctx context.Context, flushCollectionCompaction *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error) } func (s *Coordinator) ResetState(ctx context.Context) error { - return s.meta.ResetState(ctx) + return s.catalog.ResetState(ctx) } func (s *Coordinator) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error) { - database, err := s.meta.CreateDatabase(ctx, createDatabase) + database, err := s.catalog.CreateDatabase(ctx, createDatabase, createDatabase.Ts) if err != nil { return nil, err } @@ -44,7 +47,7 @@ func (s *Coordinator) CreateDatabase(ctx context.Context, createDatabase *model. } func (s *Coordinator) GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error) { - database, err := s.meta.GetDatabase(ctx, getDatabase) + database, err := s.catalog.GetDatabases(ctx, getDatabase, getDatabase.Ts) if err != nil { return nil, err } @@ -52,7 +55,7 @@ func (s *Coordinator) GetDatabase(ctx context.Context, getDatabase *model.GetDat } func (s *Coordinator) CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) { - tenant, err := s.meta.CreateTenant(ctx, createTenant) + tenant, err := s.catalog.CreateTenant(ctx, createTenant, createTenant.Ts) if err != nil { return nil, err } @@ -60,7 +63,7 @@ func (s *Coordinator) CreateTenant(ctx context.Context, createTenant *model.Crea } func (s *Coordinator) GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error) { - tenant, err := s.meta.GetTenant(ctx, getTenant) + tenant, err := s.catalog.GetTenants(ctx, getTenant, getTenant.Ts) if err != nil { return nil, err } @@ -68,68 +71,53 @@ func (s *Coordinator) GetTenant(ctx context.Context, getTenant *model.GetTenant) } func (s *Coordinator) CreateCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error) { - collectionTopic, err := s.assignCollection(createCollection.ID) - if err != nil { - return nil, err - } - createCollection.Topic = collectionTopic - log.Info("apis create collection", zap.Any("collection", createCollection)) - collection, err := s.meta.AddCollection(ctx, createCollection) + log.Info("create collection", zap.Any("createCollection", createCollection)) + collection, err := s.catalog.CreateCollection(ctx, createCollection, createCollection.Ts) if err != nil { return nil, err } return collection, nil } -func (s *Coordinator) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*model.Collection, error) { - return s.meta.GetCollections(ctx, collectionID, collectionName, collectionTopic, tenantID, databaseName) +func (s *Coordinator) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, tenantID string, databaseName string) ([]*model.Collection, error) { + return s.catalog.GetCollections(ctx, collectionID, collectionName, tenantID, databaseName) } func (s *Coordinator) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error { - return s.meta.DeleteCollection(ctx, deleteCollection) + return s.catalog.DeleteCollection(ctx, deleteCollection) } func (s *Coordinator) UpdateCollection(ctx context.Context, collection *model.UpdateCollection) (*model.Collection, error) { - return s.meta.UpdateCollection(ctx, collection) + return s.catalog.UpdateCollection(ctx, collection, collection.Ts) } func (s *Coordinator) CreateSegment(ctx context.Context, segment *model.CreateSegment) error { if err := verifyCreateSegment(segment); err != nil { return err } - err := s.meta.AddSegment(ctx, segment) + _, err := s.catalog.CreateSegment(ctx, segment, segment.Ts) if err != nil { return err } return nil } -func (s *Coordinator) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*model.Segment, error) { - return s.meta.GetSegments(ctx, segmentID, segmentType, scope, topic, collectionID) +func (s *Coordinator) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*model.Segment, error) { + return s.catalog.GetSegments(ctx, segmentID, segmentType, scope, collectionID) } func (s *Coordinator) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error { - return s.meta.DeleteSegment(ctx, segmentID) + return s.catalog.DeleteSegment(ctx, segmentID) } func (s *Coordinator) UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment) (*model.Segment, error) { - segment, err := s.meta.UpdateSegment(ctx, updateSegment) + segment, err := s.catalog.UpdateSegment(ctx, updateSegment, updateSegment.Ts) if err != nil { return nil, err } return segment, nil } -func verifyCreateCollection(collection *model.CreateCollection) error { - if collection.ID.String() == "" { - return errors.New("collection ID cannot be empty") - } - if err := verifyCollectionMetadata(collection.Metadata); err != nil { - return err - } - return nil -} - func verifyCollectionMetadata(metadata *model.CollectionMetadata[model.CollectionMetadataValueType]) error { if metadata == nil { return nil @@ -146,16 +134,6 @@ func verifyCollectionMetadata(metadata *model.CollectionMetadata[model.Collectio return nil } -func verifyUpdateCollection(collection *model.UpdateCollection) error { - if collection.ID.String() == "" { - return errors.New("collection ID cannot be empty") - } - if err := verifyCollectionMetadata(collection.Metadata); err != nil { - return err - } - return nil -} - func verifyCreateSegment(segment *model.CreateSegment) error { if err := verifySegmentMetadata(segment.Metadata); err != nil { return err @@ -163,13 +141,6 @@ func verifyCreateSegment(segment *model.CreateSegment) error { return nil } -func VerifyUpdateSegment(segment *model.UpdateSegment) error { - if err := verifySegmentMetadata(segment.Metadata); err != nil { - return err - } - return nil -} - func verifySegmentMetadata(metadata *model.SegmentMetadata[model.SegmentMetadataValueType]) error { if metadata == nil { return nil @@ -185,3 +156,15 @@ func verifySegmentMetadata(metadata *model.SegmentMetadata[model.SegmentMetadata } return nil } + +func (s *Coordinator) SetTenantLastCompactionTime(ctx context.Context, tenantID string, lastCompactionTime int64) error { + return s.catalog.SetTenantLastCompactionTime(ctx, tenantID, lastCompactionTime) +} + +func (s *Coordinator) GetTenantsLastCompactionTime(ctx context.Context, tenantIDs []string) ([]*dbmodel.Tenant, error) { + return s.catalog.GetTenantsLastCompactionTime(ctx, tenantIDs) +} + +func (s *Coordinator) FlushCollectionCompaction(ctx context.Context, flushCollectionCompaction *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error) { + return s.catalog.FlushCollectionCompaction(ctx, flushCollectionCompaction) +} diff --git a/go/pkg/coordinator/apis_test.go b/go/pkg/coordinator/apis_test.go new file mode 100644 index 00000000000..3b05931f6e5 --- /dev/null +++ b/go/pkg/coordinator/apis_test.go @@ -0,0 +1,865 @@ +package coordinator + +import ( + "context" + "sort" + "strconv" + "testing" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "gorm.io/gorm" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/google/uuid" + "pgregory.net/rapid" +) + +type APIsTestSuite struct { + suite.Suite + db *gorm.DB + collectionId1 types.UniqueID + collectionId2 types.UniqueID + records [][]byte + tenantName string + databaseName string + databaseId string + sampleCollections []*model.Collection + coordinator *Coordinator +} + +func (suite *APIsTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() +} + +func (suite *APIsTestSuite) SetupTest() { + log.Info("setup test") + suite.tenantName = "tenant_" + suite.T().Name() + suite.databaseName = "database_" + suite.T().Name() + DbId, err := dao.CreateTestTenantAndDatabase(suite.db, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.databaseId = DbId + suite.sampleCollections = SampleCollections(suite.tenantName, suite.databaseName) + for index, collection := range suite.sampleCollections { + collection.ID = types.NewUniqueID() + collection.Name = "collection_" + suite.T().Name() + strconv.Itoa(index) + } + ctx := context.Background() + c, err := NewCoordinator(ctx, suite.db, nil, nil) + if err != nil { + suite.T().Fatalf("error creating coordinator: %v", err) + } + suite.coordinator = c + for _, collection := range suite.sampleCollections { + _, errCollectionCreation := c.CreateCollection(ctx, &model.CreateCollection{ + ID: collection.ID, + Name: collection.Name, + Metadata: collection.Metadata, + Dimension: collection.Dimension, + TenantID: collection.TenantID, + DatabaseName: collection.DatabaseName, + }) + suite.NoError(errCollectionCreation) + } +} + +func (suite *APIsTestSuite) TearDownTest() { + log.Info("teardown test") + err := dao.CleanUpTestDatabase(suite.db, suite.tenantName, suite.databaseName) + suite.NoError(err) + err = dao.CleanUpTestTenant(suite.db, suite.tenantName) + suite.NoError(err) +} + +// TODO: This is not complete yet. We need to add more tests for the other APIs. +// We will deprecate the example based tests once we have enough tests here. +func testCollection(t *rapid.T) { + db := dbcore.ConfigDatabaseForTesting() + ctx := context.Background() + c, err := NewCoordinator(ctx, db, nil, nil) + if err != nil { + t.Fatalf("error creating coordinator: %v", err) + } + t.Repeat(map[string]func(*rapid.T){ + "create_collection": func(t *rapid.T) { + stringValue := generateCollectionStringMetadataValue(t) + intValue := generateCollectionInt64MetadataValue(t) + floatValue := generateCollectionFloat64MetadataValue(t) + + metadata := model.NewCollectionMetadata[model.CollectionMetadataValueType]() + metadata.Add("string_value", stringValue) + metadata.Add("int_value", intValue) + metadata.Add("float_value", floatValue) + + collection := rapid.Custom[*model.CreateCollection](func(t *rapid.T) *model.CreateCollection { + return &model.CreateCollection{ + ID: types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "collection_id")), + Name: rapid.String().Draw(t, "collection_name"), + Metadata: nil, + } + }).Draw(t, "collection") + + _, err := c.CreateCollection(ctx, collection) + if err != nil { + if err == common.ErrCollectionNameEmpty && collection.Name == "" { + t.Logf("expected error for empty collection name") + } else { + t.Fatalf("error creating collection: %v", err) + } + } + if err == nil { + // verify the correctness + collectionList, err := c.GetCollections(ctx, collection.ID, nil, common.DefaultTenant, common.DefaultDatabase) + if err != nil { + t.Fatalf("error getting collections: %v", err) + } + if len(collectionList) != 1 { + t.Fatalf("More than 1 collection with the same collection id") + } + for _, collectionReturned := range collectionList { + if collection.ID != collectionReturned.ID { + t.Fatalf("collection id is the right value") + } + } + } + }, + }) +} + +func testSegment(t *rapid.T) { + db := dbcore.ConfigDatabaseForTesting() + ctx := context.Background() + c, err := NewCoordinator(ctx, db, nil, nil) + if err != nil { + t.Fatalf("error creating coordinator: %v", err) + } + + stringValue := generateSegmentStringMetadataValue(t) + intValue := generateSegmentInt64MetadataValue(t) + floatValue := generateSegmentFloat64MetadataValue(t) + + metadata := model.NewSegmentMetadata[model.SegmentMetadataValueType]() + metadata.Set("string_value", stringValue) + metadata.Set("int_value", intValue) + metadata.Set("float_value", floatValue) + + t.Repeat(map[string]func(*rapid.T){ + "create_segment": func(t *rapid.T) { + segment := rapid.Custom[*model.CreateSegment](func(t *rapid.T) *model.CreateSegment { + return &model.CreateSegment{ + ID: types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "segment_id")), + Type: "test-segment-type", + Scope: "test-segment-scope", + Metadata: nil, + CollectionID: types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "collection_id")), + } + }).Draw(t, "segment") + + err := c.CreateSegment(ctx, segment) + if err != nil { + t.Fatalf("error creating segment: %v", err) + } + }, + }) +} + +func generateCollectionStringMetadataValue(t *rapid.T) model.CollectionMetadataValueType { + return &model.CollectionMetadataValueStringType{ + Value: rapid.String().Draw(t, "string_value"), + } +} + +func generateCollectionInt64MetadataValue(t *rapid.T) model.CollectionMetadataValueType { + return &model.CollectionMetadataValueInt64Type{ + Value: rapid.Int64().Draw(t, "int_value"), + } +} + +func generateCollectionFloat64MetadataValue(t *rapid.T) model.CollectionMetadataValueType { + return &model.CollectionMetadataValueFloat64Type{ + Value: rapid.Float64().Draw(t, "float_value"), + } +} + +func generateSegmentStringMetadataValue(t *rapid.T) model.SegmentMetadataValueType { + return &model.SegmentMetadataValueStringType{ + Value: rapid.String().Draw(t, "string_value"), + } +} + +func generateSegmentInt64MetadataValue(t *rapid.T) model.SegmentMetadataValueType { + return &model.SegmentMetadataValueInt64Type{ + Value: rapid.Int64().Draw(t, "int_value"), + } +} + +func generateSegmentFloat64MetadataValue(t *rapid.T) model.SegmentMetadataValueType { + return &model.SegmentMetadataValueFloat64Type{ + Value: rapid.Float64().Draw(t, "float_value"), + } +} + +func TestAPIs(t *testing.T) { + // rapid.Check(t, testCollection) + // rapid.Check(t, testSegment) +} + +func SampleCollections(tenantID string, databaseName string) []*model.Collection { + dimension := int32(128) + metadata1 := model.NewCollectionMetadata[model.CollectionMetadataValueType]() + metadata1.Add("test_str", &model.CollectionMetadataValueStringType{Value: "str1"}) + metadata1.Add("test_int", &model.CollectionMetadataValueInt64Type{Value: 1}) + metadata1.Add("test_float", &model.CollectionMetadataValueFloat64Type{Value: 1.3}) + + metadata2 := model.NewCollectionMetadata[model.CollectionMetadataValueType]() + metadata2.Add("test_str", &model.CollectionMetadataValueStringType{Value: "str2"}) + metadata2.Add("test_int", &model.CollectionMetadataValueInt64Type{Value: 2}) + metadata2.Add("test_float", &model.CollectionMetadataValueFloat64Type{Value: 2.3}) + + metadata3 := model.NewCollectionMetadata[model.CollectionMetadataValueType]() + metadata3.Add("test_str", &model.CollectionMetadataValueStringType{Value: "str3"}) + metadata3.Add("test_int", &model.CollectionMetadataValueInt64Type{Value: 3}) + metadata3.Add("test_float", &model.CollectionMetadataValueFloat64Type{Value: 3.3}) + sampleCollections := []*model.Collection{ + { + ID: types.MustParse("93ffe3ec-0107-48d4-8695-51f978c509dc"), + Name: "test_collection_1", + Metadata: metadata1, + Dimension: &dimension, + TenantID: tenantID, + DatabaseName: databaseName, + }, + { + ID: types.MustParse("f444f1d7-d06c-4357-ac22-5a4a1f92d761"), + Name: "test_collection_2", + Metadata: metadata2, + Dimension: nil, + TenantID: tenantID, + DatabaseName: databaseName, + }, + { + ID: types.MustParse("43babc1a-e403-4a50-91a9-16621ba29ab0"), + Name: "test_collection_3", + Metadata: metadata3, + Dimension: nil, + TenantID: tenantID, + DatabaseName: databaseName, + }, + } + return sampleCollections +} + +func (suite *APIsTestSuite) TestCreateGetDeleteCollections() { + ctx := context.Background() + results, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + + sort.Slice(results, func(i, j int) bool { + return results[i].Name < results[j].Name + }) + suite.Equal(suite.sampleCollections, results) + + // Duplicate create fails + _, err = suite.coordinator.CreateCollection(ctx, &model.CreateCollection{ + ID: suite.sampleCollections[0].ID, + Name: suite.sampleCollections[0].Name, + TenantID: suite.tenantName, + DatabaseName: suite.databaseName, + }) + suite.Error(err) + + // Find by name + for _, collection := range suite.sampleCollections { + result, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), &collection.Name, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal([]*model.Collection{collection}, result) + } + + // Find by id + for _, collection := range suite.sampleCollections { + result, err := suite.coordinator.GetCollections(ctx, collection.ID, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal([]*model.Collection{collection}, result) + } + + // Delete + c1 := suite.sampleCollections[0] + deleteCollection := &model.DeleteCollection{ + ID: c1.ID, + DatabaseName: suite.databaseName, + TenantID: suite.tenantName, + } + err = suite.coordinator.DeleteCollection(ctx, deleteCollection) + suite.NoError(err) + + results, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + + suite.NotContains(results, c1) + suite.Len(results, len(suite.sampleCollections)-1) + suite.ElementsMatch(results, suite.sampleCollections[1:]) + byIDResult, err := suite.coordinator.GetCollections(ctx, c1.ID, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Empty(byIDResult) + + // Duplicate delete throws an exception + err = suite.coordinator.DeleteCollection(ctx, deleteCollection) + suite.Error(err) +} + +func (suite *APIsTestSuite) TestUpdateCollections() { + ctx := context.Background() + coll := &model.Collection{ + Name: suite.sampleCollections[0].Name, + ID: suite.sampleCollections[0].ID, + Metadata: suite.sampleCollections[0].Metadata, + Dimension: suite.sampleCollections[0].Dimension, + TenantID: suite.sampleCollections[0].TenantID, + DatabaseName: suite.sampleCollections[0].DatabaseName, + } + + // Update name + coll.Name = "new_name" + result, err := suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Name: &coll.Name}) + suite.NoError(err) + suite.Equal(coll, result) + resultList, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), &coll.Name, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal([]*model.Collection{coll}, resultList) + + // Update dimension + newDimension := int32(128) + coll.Dimension = &newDimension + result, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Dimension: coll.Dimension}) + suite.NoError(err) + suite.Equal(coll, result) + resultList, err = suite.coordinator.GetCollections(ctx, coll.ID, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal([]*model.Collection{coll}, resultList) + + // Reset the metadata + newMetadata := model.NewCollectionMetadata[model.CollectionMetadataValueType]() + newMetadata.Add("test_str2", &model.CollectionMetadataValueStringType{Value: "str2"}) + coll.Metadata = newMetadata + result, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Metadata: coll.Metadata}) + suite.NoError(err) + suite.Equal(coll, result) + resultList, err = suite.coordinator.GetCollections(ctx, coll.ID, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal([]*model.Collection{coll}, resultList) + + // Delete all metadata keys + coll.Metadata = nil + result, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Metadata: coll.Metadata, ResetMetadata: true}) + suite.NoError(err) + suite.Equal(coll, result) + resultList, err = suite.coordinator.GetCollections(ctx, coll.ID, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal([]*model.Collection{coll}, resultList) +} + +func (suite *APIsTestSuite) TestCreateUpdateWithDatabase() { + ctx := context.Background() + newDatabaseName := "test_apis_CreateUpdateWithDatabase" + newDatabaseId := uuid.New().String() + _, err := suite.coordinator.CreateDatabase(ctx, &model.CreateDatabase{ + ID: newDatabaseId, + Name: newDatabaseName, + Tenant: suite.tenantName, + }) + suite.NoError(err) + + suite.sampleCollections[0].ID = types.NewUniqueID() + suite.sampleCollections[0].Name = suite.sampleCollections[0].Name + "1" + _, err = suite.coordinator.CreateCollection(ctx, &model.CreateCollection{ + ID: suite.sampleCollections[0].ID, + Name: suite.sampleCollections[0].Name, + Metadata: suite.sampleCollections[0].Metadata, + Dimension: suite.sampleCollections[0].Dimension, + TenantID: suite.sampleCollections[0].TenantID, + DatabaseName: newDatabaseName, + }) + suite.NoError(err) + newName1 := "new_name_1" + _, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ + ID: suite.sampleCollections[1].ID, + Name: &newName1, + }) + suite.NoError(err) + result, err := suite.coordinator.GetCollections(ctx, suite.sampleCollections[1].ID, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Len(result, 1) + suite.Equal(newName1, result[0].Name) + + newName0 := "new_name_0" + _, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ + ID: suite.sampleCollections[0].ID, + Name: &newName0, + }) + suite.NoError(err) + //suite.Equal(newName0, collection.Name) + result, err = suite.coordinator.GetCollections(ctx, suite.sampleCollections[0].ID, nil, suite.tenantName, newDatabaseName) + suite.NoError(err) + suite.Len(result, 1) + suite.Equal(newName0, result[0].Name) + + // clean up + err = dao.CleanUpTestDatabase(suite.db, suite.tenantName, newDatabaseName) + suite.NoError(err) +} + +func (suite *APIsTestSuite) TestGetMultipleWithDatabase() { + newDatabaseName := "test_apis_GetMultipleWithDatabase" + ctx := context.Background() + + newDatabaseId := uuid.New().String() + _, err := suite.coordinator.CreateDatabase(ctx, &model.CreateDatabase{ + ID: newDatabaseId, + Name: newDatabaseName, + Tenant: suite.tenantName, + }) + suite.NoError(err) + + for index, collection := range suite.sampleCollections { + collection.ID = types.NewUniqueID() + collection.Name = collection.Name + "1" + collection.TenantID = suite.tenantName + collection.DatabaseName = newDatabaseName + _, err := suite.coordinator.CreateCollection(ctx, &model.CreateCollection{ + ID: collection.ID, + Name: collection.Name, + Metadata: collection.Metadata, + Dimension: collection.Dimension, + TenantID: collection.TenantID, + DatabaseName: collection.DatabaseName, + }) + suite.NoError(err) + suite.sampleCollections[index] = collection + } + result, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, newDatabaseName) + suite.NoError(err) + suite.Equal(len(suite.sampleCollections), len(result)) + sort.Slice(result, func(i, j int) bool { + return result[i].Name < result[j].Name + }) + suite.Equal(suite.sampleCollections, result) + + result, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Equal(len(suite.sampleCollections), len(result)) + + // clean up + err = dao.CleanUpTestDatabase(suite.db, suite.tenantName, newDatabaseName) + suite.NoError(err) +} + +func (suite *APIsTestSuite) TestCreateDatabaseWithTenants() { + ctx := context.Background() + + // Create a new tenant + newTenantName := "tenant1" + _, err := suite.coordinator.CreateTenant(ctx, &model.CreateTenant{ + Name: newTenantName, + }) + suite.NoError(err) + + // Create tenant that already exits and expect an error + _, err = suite.coordinator.CreateTenant(ctx, &model.CreateTenant{ + Name: newTenantName, + }) + suite.Error(err) + + // Create tenant that already exits and expect an error + _, err = suite.coordinator.CreateTenant(ctx, &model.CreateTenant{ + Name: suite.tenantName, + }) + suite.Error(err) + + // Create a new database within this tenant and also in the default tenant + newDatabaseName := "test_apis_CreateDatabaseWithTenants" + _, err = suite.coordinator.CreateDatabase(ctx, &model.CreateDatabase{ + ID: types.MustParse("33333333-d7d7-413b-92e1-731098a6e492").String(), + Name: newDatabaseName, + Tenant: newTenantName, + }) + suite.NoError(err) + + _, err = suite.coordinator.CreateDatabase(ctx, &model.CreateDatabase{ + ID: types.MustParse("44444444-d7d7-413b-92e1-731098a6e492").String(), + Name: newDatabaseName, + Tenant: suite.tenantName, + }) + suite.NoError(err) + + // Create a new collection in the new tenant + suite.sampleCollections[0].ID = types.NewUniqueID() + suite.sampleCollections[0].Name = suite.sampleCollections[0].Name + "1" + _, err = suite.coordinator.CreateCollection(ctx, &model.CreateCollection{ + ID: suite.sampleCollections[0].ID, + Name: suite.sampleCollections[0].Name, + Metadata: suite.sampleCollections[0].Metadata, + Dimension: suite.sampleCollections[0].Dimension, + TenantID: newTenantName, + DatabaseName: newDatabaseName, + }) + suite.NoError(err) + + // Create a new collection in the default tenant + suite.sampleCollections[1].ID = types.NewUniqueID() + suite.sampleCollections[1].Name = suite.sampleCollections[1].Name + "2" + _, err = suite.coordinator.CreateCollection(ctx, &model.CreateCollection{ + ID: suite.sampleCollections[1].ID, + Name: suite.sampleCollections[1].Name, + Metadata: suite.sampleCollections[1].Metadata, + Dimension: suite.sampleCollections[1].Dimension, + TenantID: suite.tenantName, + DatabaseName: newDatabaseName, + }) + suite.NoError(err) + + // Check that both tenants have the correct collections + expected := []*model.Collection{suite.sampleCollections[0]} + expected[0].TenantID = newTenantName + expected[0].DatabaseName = newDatabaseName + result, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, newTenantName, newDatabaseName) + suite.NoError(err) + suite.Len(result, 1) + suite.Equal(expected[0], result[0]) + + expected = []*model.Collection{suite.sampleCollections[1]} + expected[0].TenantID = suite.tenantName + expected[0].DatabaseName = newDatabaseName + result, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, newDatabaseName) + suite.NoError(err) + suite.Len(result, 1) + suite.Equal(expected[0], result[0]) + + // A new tenant DOES NOT have a default database. This does not error, instead 0 + // results are returned + result, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, newTenantName, suite.databaseName) + suite.NoError(err) + suite.Nil(result) + + // clean up + err = dao.CleanUpTestTenant(suite.db, newTenantName) + suite.NoError(err) + err = dao.CleanUpTestDatabase(suite.db, suite.tenantName, newDatabaseName) + suite.NoError(err) +} + +func (suite *APIsTestSuite) TestCreateGetDeleteTenants() { + ctx := context.Background() + + // Create a new tenant + newTenantName := "tenant1" + _, err := suite.coordinator.CreateTenant(ctx, &model.CreateTenant{ + Name: newTenantName, + }) + suite.NoError(err) + + // Create tenant that already exits and expect an error + _, err = suite.coordinator.CreateTenant(ctx, &model.CreateTenant{ + Name: newTenantName, + }) + suite.Error(err) + + // Create tenant that already exits and expect an error + _, err = suite.coordinator.CreateTenant(ctx, &model.CreateTenant{ + Name: suite.tenantName, + }) + suite.Error(err) + + // Get the tenant and check that it exists + result, err := suite.coordinator.GetTenant(ctx, &model.GetTenant{Name: newTenantName}) + suite.NoError(err) + suite.Equal(newTenantName, result.Name) + + // Get a tenant that does not exist and expect an error + _, err = suite.coordinator.GetTenant(ctx, &model.GetTenant{Name: "tenant2"}) + suite.Error(err) + + // Create a new database within this tenant + newDatabaseName := "test_apis_CreateGetDeleteTenants" + _, err = suite.coordinator.CreateDatabase(ctx, &model.CreateDatabase{ + ID: types.MustParse("33333333-d7d7-413b-92e1-731098a6e492").String(), + Name: newDatabaseName, + Tenant: newTenantName, + }) + suite.NoError(err) + + // Get the database and check that it exists + databaseResult, err := suite.coordinator.GetDatabase(ctx, &model.GetDatabase{ + Name: newDatabaseName, + Tenant: newTenantName, + }) + suite.NoError(err) + suite.Equal(newDatabaseName, databaseResult.Name) + suite.Equal(newTenantName, databaseResult.Tenant) + + // Get a database that does not exist in a tenant that does exist and expect an error + _, err = suite.coordinator.GetDatabase(ctx, &model.GetDatabase{ + Name: "new_database1", + Tenant: newTenantName, + }) + suite.Error(err) + + // Get a database that does not exist in a tenant that does not exist and expect an + // error + _, err = suite.coordinator.GetDatabase(ctx, &model.GetDatabase{ + Name: "new_database1", + Tenant: "tenant2", + }) + suite.Error(err) + + // clean up + err = dao.CleanUpTestTenant(suite.db, newTenantName) + suite.NoError(err) + err = dao.CleanUpTestDatabase(suite.db, suite.tenantName, newDatabaseName) + suite.NoError(err) +} + +func SampleSegments(sampleCollections []*model.Collection) []*model.Segment { + metadata1 := model.NewSegmentMetadata[model.SegmentMetadataValueType]() + metadata1.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str1"}) + metadata1.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 1}) + metadata1.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 1.3}) + + metadata2 := model.NewSegmentMetadata[model.SegmentMetadataValueType]() + metadata2.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str2"}) + metadata2.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 2}) + metadata2.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 2.3}) + + metadata3 := model.NewSegmentMetadata[model.SegmentMetadataValueType]() + metadata3.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str3"}) + metadata3.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 3}) + metadata3.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 3.3}) + + sampleSegments := []*model.Segment{ + { + ID: types.MustParse("00000000-d7d7-413b-92e1-731098a6e492"), + Type: "test_type_a", + Scope: "VECTOR", + CollectionID: sampleCollections[0].ID, + Metadata: metadata1, + FilePaths: map[string][]string{}, + }, + { + ID: types.MustParse("11111111-d7d7-413b-92e1-731098a6e492"), + Type: "test_type_b", + Scope: "VECTOR", + CollectionID: sampleCollections[1].ID, + Metadata: metadata2, + FilePaths: map[string][]string{}, + }, + { + ID: types.MustParse("22222222-d7d7-413b-92e1-731098a6e492"), + Type: "test_type_b", + Scope: "METADATA", + CollectionID: types.NilUniqueID(), + Metadata: metadata3, // This segment is not assigned to any collection + FilePaths: map[string][]string{}, + }, + } + return sampleSegments +} + +func (suite *APIsTestSuite) TestCreateGetDeleteSegments() { + ctx := context.Background() + c := suite.coordinator + + sampleSegments := SampleSegments(suite.sampleCollections) + for _, segment := range sampleSegments { + errSegmentCreation := c.CreateSegment(ctx, &model.CreateSegment{ + ID: segment.ID, + Type: segment.Type, + Scope: segment.Scope, + CollectionID: segment.CollectionID, + Metadata: segment.Metadata, + }) + suite.NoError(errSegmentCreation) + } + + var results []*model.Segment + for _, segment := range sampleSegments { + result, err := c.GetSegments(ctx, segment.ID, nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal([]*model.Segment{segment}, result) + results = append(results, result...) + } + sort.Slice(results, func(i, j int) bool { + return results[i].ID.String() < results[j].ID.String() + }) + suite.Equal(sampleSegments, results) + + // Duplicate create fails + err := c.CreateSegment(ctx, &model.CreateSegment{ + ID: sampleSegments[0].ID, + Type: sampleSegments[0].Type, + Scope: sampleSegments[0].Scope, + CollectionID: sampleSegments[0].CollectionID, + Metadata: sampleSegments[0].Metadata, + }) + suite.Error(err) + + // Find by id + for _, segment := range sampleSegments { + result, err := c.GetSegments(ctx, segment.ID, nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal([]*model.Segment{segment}, result) + } + + // Find by type + testTypeA := "test_type_a" + result, err := c.GetSegments(ctx, types.NilUniqueID(), &testTypeA, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal(sampleSegments[:1], result) + + testTypeB := "test_type_b" + result, err = c.GetSegments(ctx, types.NilUniqueID(), &testTypeB, nil, types.NilUniqueID()) + suite.NoError(err) + suite.ElementsMatch(sampleSegments[1:], result) + + // Find by collection ID + result, err = c.GetSegments(ctx, types.NilUniqueID(), nil, nil, suite.sampleCollections[0].ID) + suite.NoError(err) + suite.Equal(sampleSegments[:1], result) + + // Find by type and collection ID (positive case) + result, err = c.GetSegments(ctx, types.NilUniqueID(), &testTypeA, nil, suite.sampleCollections[0].ID) + suite.NoError(err) + suite.Equal(sampleSegments[:1], result) + + // Find by type and collection ID (negative case) + result, err = c.GetSegments(ctx, types.NilUniqueID(), &testTypeB, nil, suite.sampleCollections[0].ID) + suite.NoError(err) + suite.Empty(result) + + // Delete + s1 := sampleSegments[0] + err = c.DeleteSegment(ctx, s1.ID) + suite.NoError(err) + + results, err = c.GetSegments(ctx, types.NilUniqueID(), nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.NotContains(results, s1) + suite.Len(results, len(sampleSegments)-1) + suite.ElementsMatch(results, sampleSegments[1:]) + + // Duplicate delete throws an exception + err = c.DeleteSegment(ctx, s1.ID) + suite.Error(err) + + // clean up segments + for _, segment := range sampleSegments { + _ = c.DeleteSegment(ctx, segment.ID) + } +} + +func (suite *APIsTestSuite) TestUpdateSegment() { + metadata := model.NewSegmentMetadata[model.SegmentMetadataValueType]() + metadata.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str1"}) + metadata.Set("test_int", &model.SegmentMetadataValueInt64Type{Value: 1}) + metadata.Set("test_float", &model.SegmentMetadataValueFloat64Type{Value: 1.3}) + + segment := &model.Segment{ + ID: types.UniqueID(uuid.New()), + Type: "test_type_a", + Scope: "VECTOR", + CollectionID: suite.sampleCollections[0].ID, + Metadata: metadata, + FilePaths: map[string][]string{}, + } + + ctx := context.Background() + errSegmentCreation := suite.coordinator.CreateSegment(ctx, &model.CreateSegment{ + ID: segment.ID, + Type: segment.Type, + Scope: segment.Scope, + CollectionID: segment.CollectionID, + Metadata: segment.Metadata, + }) + suite.NoError(errSegmentCreation) + + collectionID := segment.CollectionID.String() + + // TODO: revisit why we need this + // Update collection to new value + //segment.CollectionID = sampleCollections[1].ID + //newCollecionID := segment.CollectionID.String() + //c.UpdateSegment(ctx, &model.UpdateSegment{ + // ID: segment.ID, + // Collection: &newCollecionID, + //}) + //result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) + //assert.NoError(t, err) + //assert.Equal(t, []*model.Segment{segment}, result) + + // Update collection to None + //segment.CollectionID = types.NilUniqueID() + //c.UpdateSegment(ctx, &model.UpdateSegment{ + // ID: segment.ID, + // Collection: nil, + // ResetCollection: true, + //}) + //result, err = c.GetSegments(ctx, segment.ID, nil, nil, nil, types.NilUniqueID()) + //assert.NoError(t, err) + //assert.Equal(t, []*model.Segment{segment}, result) + + // Add a new metadata key + segment.Metadata.Set("test_str2", &model.SegmentMetadataValueStringType{Value: "str2"}) + _, err := suite.coordinator.UpdateSegment(ctx, &model.UpdateSegment{ + Collection: &collectionID, + ID: segment.ID, + Metadata: segment.Metadata}) + suite.NoError(err) + result, err := suite.coordinator.GetSegments(ctx, segment.ID, nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal([]*model.Segment{segment}, result) + + // Update a metadata key + segment.Metadata.Set("test_str", &model.SegmentMetadataValueStringType{Value: "str3"}) + _, err = suite.coordinator.UpdateSegment(ctx, &model.UpdateSegment{ + Collection: &collectionID, + ID: segment.ID, + Metadata: segment.Metadata}) + suite.NoError(err) + result, err = suite.coordinator.GetSegments(ctx, segment.ID, nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal([]*model.Segment{segment}, result) + + // Delete a metadata key + segment.Metadata.Remove("test_str") + newMetadata := model.NewSegmentMetadata[model.SegmentMetadataValueType]() + newMetadata.Set("test_str", nil) + _, err = suite.coordinator.UpdateSegment(ctx, &model.UpdateSegment{ + Collection: &collectionID, + ID: segment.ID, + Metadata: newMetadata}) + suite.NoError(err) + result, err = suite.coordinator.GetSegments(ctx, segment.ID, nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal([]*model.Segment{segment}, result) + + // Delete all metadata keys + segment.Metadata = nil + _, err = suite.coordinator.UpdateSegment(ctx, &model.UpdateSegment{ + Collection: &collectionID, + ID: segment.ID, + Metadata: segment.Metadata, + ResetMetadata: true}, + ) + suite.NoError(err) + result, err = suite.coordinator.GetSegments(ctx, segment.ID, nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Equal([]*model.Segment{segment}, result) +} + +func TestAPIsTestSuite(t *testing.T) { + testSuite := new(APIsTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/pkg/coordinator/coordinator.go b/go/pkg/coordinator/coordinator.go new file mode 100644 index 00000000000..abef2c29ca7 --- /dev/null +++ b/go/pkg/coordinator/coordinator.go @@ -0,0 +1,56 @@ +package coordinator + +import ( + "context" + "log" + + "github.com/chroma-core/chroma/go/pkg/metastore" + "github.com/chroma-core/chroma/go/pkg/metastore/coordinator" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/notification" + "gorm.io/gorm" +) + +var _ ICoordinator = (*Coordinator)(nil) + +// Coordinator is the implemenation of ICoordinator. It is the top level component. +// Currently, it only has the system catalog related APIs and will be extended to +// support other functionalities such as membership managed and propagation. +type Coordinator struct { + ctx context.Context + notificationProcessor notification.NotificationProcessor + catalog metastore.Catalog +} + +func NewCoordinator(ctx context.Context, db *gorm.DB, notificationStore notification.NotificationStore, notifier notification.Notifier) (*Coordinator, error) { + s := &Coordinator{ + ctx: ctx, + } + + notificationProcessor := notification.NewSimpleNotificationProcessor(ctx, notificationStore, notifier) + s.notificationProcessor = notificationProcessor + + // catalog + txnImpl := dbcore.NewTxImpl() + metaDomain := dao.NewMetaDomain() + s.catalog = coordinator.NewTableCatalogWithNotification(txnImpl, metaDomain, notificationStore) + return s, nil +} + +func (s *Coordinator) Start() error { + err := s.notificationProcessor.Start() + if err != nil { + log.Printf("Failed to start notification processor: %v", err) + return err + } + return nil +} + +func (s *Coordinator) Stop() error { + err := s.notificationProcessor.Stop() + if err != nil { + log.Printf("Failed to stop notification processor: %v", err) + } + return nil +} diff --git a/go/coordinator/internal/grpccoordinator/collection_service.go b/go/pkg/coordinator/grpc/collection_service.go similarity index 70% rename from go/coordinator/internal/grpccoordinator/collection_service.go rename to go/pkg/coordinator/grpc/collection_service.go index faaf6b4dbf9..ad724a1ca41 100644 --- a/go/coordinator/internal/grpccoordinator/collection_service.go +++ b/go/pkg/coordinator/grpc/collection_service.go @@ -1,12 +1,16 @@ -package grpccoordinator +package grpc import ( "context" + "encoding/json" + "errors" - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" "google.golang.org/protobuf/types/known/emptypb" @@ -16,8 +20,9 @@ const errorCode = 500 const successCode = 200 const success = "ok" -func (s *Server) ResetState(context.Context, *emptypb.Empty) (*coordinatorpb.ChromaResponse, error) { - res := &coordinatorpb.ChromaResponse{} +func (s *Server) ResetState(context.Context, *emptypb.Empty) (*coordinatorpb.ResetStateResponse, error) { + log.Info("reset state") + res := &coordinatorpb.ResetStateResponse{} err := s.coordinator.ResetState(context.Background()) if err != nil { res.Status = failResponseWithError(err, errorCode) @@ -86,7 +91,6 @@ func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.Create return res, nil } res.Collection = convertCollectionToProto(collection) - res.Created = collection.Created res.Status = setResponseStatus(successCode) return res, nil } @@ -94,7 +98,6 @@ func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.Create func (s *Server) GetCollections(ctx context.Context, req *coordinatorpb.GetCollectionsRequest) (*coordinatorpb.GetCollectionsResponse, error) { collectionID := req.Id collectionName := req.Name - collectionTopic := req.Topic tenantID := req.Tenant databaseName := req.Database @@ -107,7 +110,7 @@ func (s *Server) GetCollections(ctx context.Context, req *coordinatorpb.GetColle return res, nil } - collections, err := s.coordinator.GetCollections(ctx, parsedCollectionID, collectionName, collectionTopic, tenantID, databaseName) + collections, err := s.coordinator.GetCollections(ctx, parsedCollectionID, collectionName, tenantID, databaseName) if err != nil { log.Error("error getting collections", zap.Error(err)) res.Status = failResponseWithError(err, errorCode) @@ -123,9 +126,9 @@ func (s *Server) GetCollections(ctx context.Context, req *coordinatorpb.GetColle return res, nil } -func (s *Server) DeleteCollection(ctx context.Context, req *coordinatorpb.DeleteCollectionRequest) (*coordinatorpb.ChromaResponse, error) { +func (s *Server) DeleteCollection(ctx context.Context, req *coordinatorpb.DeleteCollectionRequest) (*coordinatorpb.DeleteCollectionResponse, error) { collectionID := req.GetId() - res := &coordinatorpb.ChromaResponse{} + res := &coordinatorpb.DeleteCollectionResponse{} parsedCollectionID, err := types.Parse(collectionID) if err != nil { log.Error(err.Error(), zap.String("collectionpd.id", collectionID)) @@ -139,10 +142,11 @@ func (s *Server) DeleteCollection(ctx context.Context, req *coordinatorpb.Delete } err = s.coordinator.DeleteCollection(ctx, deleteCollection) if err != nil { - log.Error(err.Error(), zap.String("collectionpd.id", collectionID)) - if err == common.ErrCollectionDeleteNonExistingCollection { + if errors.Is(err, common.ErrCollectionDeleteNonExistingCollection) { + log.Error("ErrCollectionDeleteNonExistingCollection", zap.String("collectionpd.id", collectionID)) res.Status = failResponseWithError(err, 404) } else { + log.Error(err.Error(), zap.String("collectionpd.id", collectionID)) res.Status = failResponseWithError(err, errorCode) } return res, nil @@ -151,8 +155,8 @@ func (s *Server) DeleteCollection(ctx context.Context, req *coordinatorpb.Delete return res, nil } -func (s *Server) UpdateCollection(ctx context.Context, req *coordinatorpb.UpdateCollectionRequest) (*coordinatorpb.ChromaResponse, error) { - res := &coordinatorpb.ChromaResponse{} +func (s *Server) UpdateCollection(ctx context.Context, req *coordinatorpb.UpdateCollectionRequest) (*coordinatorpb.UpdateCollectionResponse, error) { + res := &coordinatorpb.UpdateCollectionResponse{} collectionID := req.Id parsedCollectionID, err := types.ToUniqueID(&collectionID) @@ -165,7 +169,6 @@ func (s *Server) UpdateCollection(ctx context.Context, req *coordinatorpb.Update updateCollection := &model.UpdateCollection{ ID: parsedCollectionID, Name: req.Name, - Topic: req.Topic, Dimension: req.Dimension, } @@ -209,6 +212,53 @@ func (s *Server) UpdateCollection(ctx context.Context, req *coordinatorpb.Update return res, nil } +func (s *Server) FlushCollectionCompaction(ctx context.Context, req *coordinatorpb.FlushCollectionCompactionRequest) (*coordinatorpb.FlushCollectionCompactionResponse, error) { + blob, err := json.Marshal(req) + if err != nil { + return nil, err + } + log.Info("flush collection compaction", zap.String("request", string(blob))) + collectionID, err := types.ToUniqueID(&req.CollectionId) + err = grpcutils.BuildErrorForUUID(collectionID, "collection", err) + if err != nil { + return nil, err + } + segmentCompactionInfo := make([]*model.FlushSegmentCompaction, 0, len(req.SegmentCompactionInfo)) + for _, flushSegmentCompaction := range req.SegmentCompactionInfo { + segmentID, err := types.ToUniqueID(&flushSegmentCompaction.SegmentId) + err = grpcutils.BuildErrorForUUID(segmentID, "segment", err) + if err != nil { + return nil, err + } + filePaths := make(map[string][]string) + for key, filePath := range flushSegmentCompaction.FilePaths { + filePaths[key] = filePath.Paths + } + segmentCompactionInfo = append(segmentCompactionInfo, &model.FlushSegmentCompaction{ + ID: segmentID, + FilePaths: filePaths, + }) + } + FlushCollectionCompaction := &model.FlushCollectionCompaction{ + ID: collectionID, + TenantID: req.TenantId, + LogPosition: req.LogPosition, + CurrentCollectionVersion: req.CollectionVersion, + FlushSegmentCompactions: segmentCompactionInfo, + } + flushCollectionInfo, err := s.coordinator.FlushCollectionCompaction(ctx, FlushCollectionCompaction) + if err != nil { + log.Error("error FlushCollectionCompaction", zap.Error(err)) + return nil, grpcutils.BuildInternalGrpcError(err.Error()) + } + res := &coordinatorpb.FlushCollectionCompactionResponse{ + CollectionId: flushCollectionInfo.ID, + CollectionVersion: flushCollectionInfo.CollectionVersion, + LastCompactionTime: flushCollectionInfo.TenantLastCompactionTime, + } + return res, nil +} + func failResponseWithError(err error, code int32) *coordinatorpb.Status { return &coordinatorpb.Status{ Reason: err.Error(), diff --git a/go/pkg/coordinator/grpc/collection_service_test.go b/go/pkg/coordinator/grpc/collection_service_test.go new file mode 100644 index 00000000000..dd4ddd75747 --- /dev/null +++ b/go/pkg/coordinator/grpc/collection_service_test.go @@ -0,0 +1,334 @@ +package grpc + +import ( + "context" + "reflect" + "strconv" + "testing" + "time" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/chroma-core/chroma/go/pkg/metastore/coordinator" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "google.golang.org/genproto/googleapis/rpc/code" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "gorm.io/gorm" + "k8s.io/apimachinery/pkg/util/rand" + "pgregory.net/rapid" +) + +type CollectionServiceTestSuite struct { + suite.Suite + catalog *coordinator.Catalog + db *gorm.DB + s *Server + tenantName string + databaseName string + databaseId string +} + +func (suite *CollectionServiceTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() + s, err := NewWithGrpcProvider(Config{ + SystemCatalogProvider: "database", + NotificationStoreProvider: "memory", + NotifierProvider: "memory", + Testing: true}, grpcutils.Default, suite.db) + if err != nil { + suite.T().Fatalf("error creating server: %v", err) + } + suite.s = s + txnImpl := dbcore.NewTxImpl() + metaDomain := dao.NewMetaDomain() + suite.catalog = coordinator.NewTableCatalogWithNotification(txnImpl, metaDomain, nil) + suite.tenantName = "tenant_" + suite.T().Name() + suite.databaseName = "database_" + suite.T().Name() + DbId, err := dao.CreateTestTenantAndDatabase(suite.db, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.databaseId = DbId +} + +func (suite *CollectionServiceTestSuite) TearDownSuite() { + log.Info("teardown suite") + err := dao.CleanUpTestDatabase(suite.db, suite.tenantName, suite.databaseName) + suite.NoError(err) + err = dao.CleanUpTestTenant(suite.db, suite.tenantName) + suite.NoError(err) +} + +// CreateCollection +// Collection created successfully are visible to ListCollections +// Collection created should have the right metadata, the metadata should be a flat map, with keys as strings and values as strings, ints, or floats +// Collection created should have the right name +// Collection created should have the right ID +// Collection created should have the right timestamp +func testCollection(t *rapid.T) { + db := dbcore.ConfigDatabaseForTesting() + s, err := NewWithGrpcProvider(Config{ + SystemCatalogProvider: "memory", + NotificationStoreProvider: "memory", + NotifierProvider: "memory", + Testing: true}, grpcutils.Default, db) + if err != nil { + t.Fatalf("error creating server: %v", err) + } + var state []*coordinatorpb.Collection + var collectionsWithErrors []*coordinatorpb.Collection + + t.Repeat(map[string]func(*rapid.T){ + "create_collection": func(t *rapid.T) { + stringValue := generateStringMetadataValue(t) + intValue := generateInt64MetadataValue(t) + floatValue := generateFloat64MetadataValue(t) + getOrCreate := false + + createCollectionRequest := rapid.Custom[*coordinatorpb.CreateCollectionRequest](func(t *rapid.T) *coordinatorpb.CreateCollectionRequest { + return &coordinatorpb.CreateCollectionRequest{ + Id: rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "collection_id"), + Name: rapid.String().Draw(t, "collection_name"), + Metadata: &coordinatorpb.UpdateMetadata{ + Metadata: map[string]*coordinatorpb.UpdateMetadataValue{ + "string_value": stringValue, + "int_value": intValue, + "float_value": floatValue, + }, + }, + GetOrCreate: &getOrCreate, + } + }).Draw(t, "create_collection_request") + + ctx := context.Background() + res, err := s.CreateCollection(ctx, createCollectionRequest) + if err != nil { + if err == common.ErrCollectionNameEmpty && createCollectionRequest.Name == "" { + t.Logf("expected error for empty collection name") + collectionsWithErrors = append(collectionsWithErrors, res.Collection) + } else { + t.Fatalf("error creating collection: %v", err) + collectionsWithErrors = append(collectionsWithErrors, res.Collection) + } + } + + getCollectionsRequest := coordinatorpb.GetCollectionsRequest{ + Id: &createCollectionRequest.Id, + } + if err == nil { + // verify the correctness + GetCollectionsResponse, err := s.GetCollections(ctx, &getCollectionsRequest) + if err != nil { + t.Fatalf("error getting collections: %v", err) + } + collectionList := GetCollectionsResponse.GetCollections() + if len(collectionList) != 1 { + t.Fatalf("More than 1 collection with the same collection id") + } + for _, collection := range collectionList { + if collection.Id != createCollectionRequest.Id { + t.Fatalf("collection id is the right value") + } + } + state = append(state, res.Collection) + } + }, + "get_collections": func(t *rapid.T) { + }, + }) +} + +func generateStringMetadataValue(t *rapid.T) *coordinatorpb.UpdateMetadataValue { + return &coordinatorpb.UpdateMetadataValue{ + Value: &coordinatorpb.UpdateMetadataValue_StringValue{ + StringValue: rapid.String().Draw(t, "string_value"), + }, + } +} + +func generateInt64MetadataValue(t *rapid.T) *coordinatorpb.UpdateMetadataValue { + return &coordinatorpb.UpdateMetadataValue{ + Value: &coordinatorpb.UpdateMetadataValue_IntValue{ + IntValue: rapid.Int64().Draw(t, "int_value"), + }, + } +} + +func generateFloat64MetadataValue(t *rapid.T) *coordinatorpb.UpdateMetadataValue { + return &coordinatorpb.UpdateMetadataValue{ + Value: &coordinatorpb.UpdateMetadataValue_FloatValue{ + FloatValue: rapid.Float64().Draw(t, "float_value"), + }, + } +} + +func TestCollection(t *testing.T) { + // rapid.Check(t, testCollection) +} + +func validateDatabase(suite *CollectionServiceTestSuite, collectionId string, collection *coordinatorpb.Collection, filePaths map[string]map[string]*coordinatorpb.FilePaths) { + getCollectionReq := coordinatorpb.GetCollectionsRequest{ + Id: &collectionId, + } + collectionsInDB, err := suite.s.GetCollections(context.Background(), &getCollectionReq) + suite.NoError(err) + suite.Len(collectionsInDB.Collections, 1) + suite.Equal(collection.Id, collection.Id) + suite.Equal(collection.Name, collection.Name) + suite.Equal(collection.LogPosition, collection.LogPosition) + suite.Equal(collection.Version, collection.Version) + + getSegmentReq := coordinatorpb.GetSegmentsRequest{ + Collection: &collectionId, + } + segments, err := suite.s.GetSegments(context.Background(), &getSegmentReq) + suite.NoError(err) + for _, segment := range segments.Segments { + suite.True(reflect.DeepEqual(filePaths[segment.Id], segment.FilePaths)) + } +} + +func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction() { + log.Info("TestServer_FlushCollectionCompaction") + // create test collection + collectionName := "collection_service_test_flush_collection_compaction" + collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId) + suite.NoError(err) + + // flush collection compaction + getSegmentReq := coordinatorpb.GetSegmentsRequest{ + Collection: &collectionID, + } + segments, err := suite.s.GetSegments(context.Background(), &getSegmentReq) + suite.NoError(err) + + flushInfo := make([]*coordinatorpb.FlushSegmentCompactionInfo, 0, len(segments.Segments)) + filePaths := make(map[string]map[string]*coordinatorpb.FilePaths, 0) + testFilePathTypes := []string{"TypeA", "TypeB", "TypeC", "TypeD"} + for _, segment := range segments.Segments { + filePaths[segment.Id] = make(map[string]*coordinatorpb.FilePaths, 0) + for i := 0; i < rand.Intn(len(testFilePathTypes)); i++ { + filePathsThisSeg := make([]string, 0) + for j := 0; j < rand.Intn(5); j++ { + filePathsThisSeg = append(filePathsThisSeg, "test_file_path_"+strconv.Itoa(j+1)) + } + filePathTypeI := rand.Intn(len(testFilePathTypes)) + filePaths[segment.Id][testFilePathTypes[filePathTypeI]] = &coordinatorpb.FilePaths{ + Paths: filePathsThisSeg, + } + } + info := &coordinatorpb.FlushSegmentCompactionInfo{ + SegmentId: segment.Id, + FilePaths: filePaths[segment.Id], + } + flushInfo = append(flushInfo, info) + } + + req := &coordinatorpb.FlushCollectionCompactionRequest{ + TenantId: suite.tenantName, + CollectionId: collectionID, + LogPosition: 10, + CollectionVersion: 0, + SegmentCompactionInfo: flushInfo, + } + response, err := suite.s.FlushCollectionCompaction(context.Background(), req) + t1 := time.Now().Unix() + suite.NoError(err) + suite.Equal(collectionID, response.CollectionId) + suite.Equal(int32(1), response.CollectionVersion) + suite.Less(int64(0), response.LastCompactionTime) + suite.LessOrEqual(response.LastCompactionTime, t1) + + // validate database + collection := &coordinatorpb.Collection{ + Id: collectionID, + LogPosition: int64(10), + Version: int32(1), + } + validateDatabase(suite, collectionID, collection, filePaths) + + // flush one segment + filePaths[segments.Segments[0].Id][testFilePathTypes[0]] = &coordinatorpb.FilePaths{ + Paths: []string{"test_file_path_1"}, + } + info := &coordinatorpb.FlushSegmentCompactionInfo{ + SegmentId: segments.Segments[0].Id, + FilePaths: filePaths[segments.Segments[0].Id], + } + req = &coordinatorpb.FlushCollectionCompactionRequest{ + TenantId: suite.tenantName, + CollectionId: collectionID, + LogPosition: 100, + CollectionVersion: 1, + SegmentCompactionInfo: []*coordinatorpb.FlushSegmentCompactionInfo{info}, + } + response, err = suite.s.FlushCollectionCompaction(context.Background(), req) + t2 := time.Now().Unix() + suite.NoError(err) + suite.Equal(collectionID, response.CollectionId) + suite.Equal(int32(2), response.CollectionVersion) + suite.LessOrEqual(t1, response.LastCompactionTime) + suite.LessOrEqual(response.LastCompactionTime, t2) + + // validate database + collection = &coordinatorpb.Collection{ + Id: collectionID, + LogPosition: int64(100), + Version: int32(2), + } + validateDatabase(suite, collectionID, collection, filePaths) + + // test invalid log position + req = &coordinatorpb.FlushCollectionCompactionRequest{ + TenantId: suite.tenantName, + CollectionId: collectionID, + LogPosition: 50, + CollectionVersion: 2, + SegmentCompactionInfo: []*coordinatorpb.FlushSegmentCompactionInfo{info}, + } + response, err = suite.s.FlushCollectionCompaction(context.Background(), req) + suite.Error(err) + suite.Equal(status.Error(codes.Code(code.Code_INTERNAL), common.ErrCollectionLogPositionStale.Error()), err) + // nothing should change in DB + validateDatabase(suite, collectionID, collection, filePaths) + + // test invalid version + req = &coordinatorpb.FlushCollectionCompactionRequest{ + TenantId: suite.tenantName, + CollectionId: collectionID, + LogPosition: 150, + CollectionVersion: 1, + SegmentCompactionInfo: []*coordinatorpb.FlushSegmentCompactionInfo{info}, + } + response, err = suite.s.FlushCollectionCompaction(context.Background(), req) + suite.Error(err) + suite.Equal(status.Error(codes.Code(code.Code_INTERNAL), common.ErrCollectionVersionStale.Error()), err) + // nothing should change in DB + validateDatabase(suite, collectionID, collection, filePaths) + + req = &coordinatorpb.FlushCollectionCompactionRequest{ + TenantId: suite.tenantName, + CollectionId: collectionID, + LogPosition: 150, + CollectionVersion: 5, + SegmentCompactionInfo: []*coordinatorpb.FlushSegmentCompactionInfo{info}, + } + response, err = suite.s.FlushCollectionCompaction(context.Background(), req) + suite.Error(err) + suite.Equal(status.Error(codes.Code(code.Code_INTERNAL), common.ErrCollectionVersionInvalid.Error()), err) + // nothing should change in DB + validateDatabase(suite, collectionID, collection, filePaths) + + // clean up + err = dao.CleanUpTestCollection(suite.db, collectionID) + suite.NoError(err) +} + +func TestCollectionServiceTestSuite(t *testing.T) { + testSuite := new(CollectionServiceTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/coordinator/internal/grpccoordinator/proto_model_convert.go b/go/pkg/coordinator/grpc/proto_model_convert.go similarity index 91% rename from go/coordinator/internal/grpccoordinator/proto_model_convert.go rename to go/pkg/coordinator/grpc/proto_model_convert.go index 18c4fd307ab..cc7fbb12fcb 100644 --- a/go/coordinator/internal/grpccoordinator/proto_model_convert.go +++ b/go/pkg/coordinator/grpc/proto_model_convert.go @@ -1,10 +1,10 @@ -package grpccoordinator +package grpc import ( - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" ) @@ -38,12 +38,13 @@ func convertCollectionToProto(collection *model.Collection) *coordinatorpb.Colle } collectionpb := &coordinatorpb.Collection{ - Id: collection.ID.String(), - Name: collection.Name, - Topic: collection.Topic, - Dimension: collection.Dimension, - Tenant: collection.TenantID, - Database: collection.DatabaseName, + Id: collection.ID.String(), + Name: collection.Name, + Dimension: collection.Dimension, + Tenant: collection.TenantID, + Database: collection.DatabaseName, + LogPosition: collection.LogPosition, + Version: collection.Version, } if collection.Metadata == nil { return collectionpb @@ -145,13 +146,19 @@ func convertSegmentToProto(segment *model.Segment) *coordinatorpb.Segment { } scope := coordinatorpb.SegmentScope_value[segment.Scope] segmentSceope := coordinatorpb.SegmentScope(scope) + filePaths := make(map[string]*coordinatorpb.FilePaths) + for t, paths := range segment.FilePaths { + filePaths[t] = &coordinatorpb.FilePaths{ + Paths: paths, + } + } segmentpb := &coordinatorpb.Segment{ Id: segment.ID.String(), Type: segment.Type, Scope: segmentSceope, - Topic: segment.Topic, Collection: nil, Metadata: nil, + FilePaths: filePaths, } collectionID := segment.CollectionID @@ -220,7 +227,6 @@ func convertSegmentToModel(segmentpb *coordinatorpb.Segment) (*model.CreateSegme ID: segmentID, Type: segmentpb.Type, Scope: segmentpb.Scope.String(), - Topic: segmentpb.Topic, CollectionID: collectionID, Metadata: metadata, }, nil diff --git a/go/coordinator/internal/grpccoordinator/proto_model_convert_test.go b/go/pkg/coordinator/grpc/proto_model_convert_test.go similarity index 93% rename from go/coordinator/internal/grpccoordinator/proto_model_convert_test.go rename to go/pkg/coordinator/grpc/proto_model_convert_test.go index 9cfa2f0632f..659addbc8d5 100644 --- a/go/coordinator/internal/grpccoordinator/proto_model_convert_test.go +++ b/go/pkg/coordinator/grpc/proto_model_convert_test.go @@ -1,11 +1,11 @@ -package grpccoordinator +package grpc import ( "testing" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/stretchr/testify/assert" ) @@ -53,7 +53,6 @@ func TestConvertCollectionToProto(t *testing.T) { collection := &model.Collection{ ID: types.NewUniqueID(), Name: "test_collection", - Topic: "test_topic", Dimension: &dimention, Metadata: &model.CollectionMetadata[model.CollectionMetadataValueType]{ Metadata: map[string]model.CollectionMetadataValueType{ @@ -67,7 +66,6 @@ func TestConvertCollectionToProto(t *testing.T) { assert.NotNil(t, collectionpb) assert.Equal(t, collection.ID.String(), collectionpb.Id) assert.Equal(t, collection.Name, collectionpb.Name) - assert.Equal(t, collection.Topic, collectionpb.Topic) assert.Equal(t, collection.Dimension, collectionpb.Dimension) assert.NotNil(t, collectionpb.Metadata) assert.Equal(t, "value1", collectionpb.Metadata.Metadata["key1"].GetStringValue()) @@ -182,20 +180,18 @@ func TestConvertSegmentToProto(t *testing.T) { assert.Nil(t, segmentpb) // Test case 2: segment is not nil - testTopic := "test_topic" segment := &model.Segment{ - ID: types.NewUniqueID(), - Type: "test_type", - Scope: "METADATA", - Topic: &testTopic, - Metadata: nil, + ID: types.NewUniqueID(), + Type: "test_type", + Scope: "METADATA", + Metadata: nil, + FilePaths: map[string][]string{}, } segmentpb = convertSegmentToProto(segment) assert.NotNil(t, segmentpb) assert.Equal(t, segment.ID.String(), segmentpb.Id) assert.Equal(t, segment.Type, segmentpb.Type) assert.Equal(t, coordinatorpb.SegmentScope_METADATA, segmentpb.Scope) - assert.Equal(t, segment.Topic, segmentpb.Topic) assert.Nil(t, segmentpb.Collection) assert.Nil(t, segmentpb.Metadata) } diff --git a/go/coordinator/internal/grpccoordinator/segment_service.go b/go/pkg/coordinator/grpc/segment_service.go similarity index 83% rename from go/coordinator/internal/grpccoordinator/segment_service.go rename to go/pkg/coordinator/grpc/segment_service.go index b2d3be5e4ff..ae3da6595eb 100644 --- a/go/coordinator/internal/grpccoordinator/segment_service.go +++ b/go/pkg/coordinator/grpc/segment_service.go @@ -1,20 +1,20 @@ -package grpccoordinator +package grpc import ( "context" - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" ) -func (s *Server) CreateSegment(ctx context.Context, req *coordinatorpb.CreateSegmentRequest) (*coordinatorpb.ChromaResponse, error) { +func (s *Server) CreateSegment(ctx context.Context, req *coordinatorpb.CreateSegmentRequest) (*coordinatorpb.CreateSegmentResponse, error) { segmentpb := req.GetSegment() - res := &coordinatorpb.ChromaResponse{} + res := &coordinatorpb.CreateSegmentResponse{} segment, err := convertSegmentToModel(segmentpb) if err != nil { @@ -43,7 +43,6 @@ func (s *Server) GetSegments(ctx context.Context, req *coordinatorpb.GetSegments segmentID := req.Id segmentType := req.Type scope := req.Scope - topic := req.Topic collectionID := req.Collection res := &coordinatorpb.GetSegmentsResponse{} @@ -67,7 +66,7 @@ func (s *Server) GetSegments(ctx context.Context, req *coordinatorpb.GetSegments scopeString := scope.String() scopeValue = &scopeString } - segments, err := s.coordinator.GetSegments(ctx, parsedSegmentID, segmentType, scopeValue, topic, parsedCollectionID) + segments, err := s.coordinator.GetSegments(ctx, parsedSegmentID, segmentType, scopeValue, parsedCollectionID) if err != nil { log.Error("get segments error", zap.Error(err)) res.Status = failResponseWithError(err, errorCode) @@ -84,9 +83,9 @@ func (s *Server) GetSegments(ctx context.Context, req *coordinatorpb.GetSegments return res, nil } -func (s *Server) DeleteSegment(ctx context.Context, req *coordinatorpb.DeleteSegmentRequest) (*coordinatorpb.ChromaResponse, error) { +func (s *Server) DeleteSegment(ctx context.Context, req *coordinatorpb.DeleteSegmentRequest) (*coordinatorpb.DeleteSegmentResponse, error) { segmentID := req.GetId() - res := &coordinatorpb.ChromaResponse{} + res := &coordinatorpb.DeleteSegmentResponse{} parsedSegmentID, err := types.Parse(segmentID) if err != nil { log.Error(err.Error(), zap.String("segment.id", segmentID)) @@ -108,20 +107,14 @@ func (s *Server) DeleteSegment(ctx context.Context, req *coordinatorpb.DeleteSeg return res, nil } -func (s *Server) UpdateSegment(ctx context.Context, req *coordinatorpb.UpdateSegmentRequest) (*coordinatorpb.ChromaResponse, error) { - res := &coordinatorpb.ChromaResponse{} +func (s *Server) UpdateSegment(ctx context.Context, req *coordinatorpb.UpdateSegmentRequest) (*coordinatorpb.UpdateSegmentResponse, error) { + res := &coordinatorpb.UpdateSegmentResponse{} updateSegment := &model.UpdateSegment{ ID: types.MustParse(req.Id), - ResetTopic: req.GetResetTopic(), ResetCollection: req.GetResetCollection(), ResetMetadata: req.GetResetMetadata(), } - topic := req.GetTopic() - if topic == "" { - updateSegment.Topic = nil - } else { - updateSegment.Topic = &topic - } + collection := req.GetCollection() if collection == "" { updateSegment.Collection = nil diff --git a/go/coordinator/internal/grpccoordinator/server.go b/go/pkg/coordinator/grpc/server.go similarity index 63% rename from go/coordinator/internal/grpccoordinator/server.go rename to go/pkg/coordinator/grpc/server.go index 4205a47153b..96b4f53ce9e 100644 --- a/go/coordinator/internal/grpccoordinator/server.go +++ b/go/pkg/coordinator/grpc/server.go @@ -1,19 +1,20 @@ -package grpccoordinator +package grpc import ( "context" "errors" "time" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/apache/pulsar-client-go/pulsar" - "github.com/chroma/chroma-coordinator/internal/coordinator" - "github.com/chroma/chroma-coordinator/internal/grpccoordinator/grpcutils" - "github.com/chroma/chroma-coordinator/internal/memberlist_manager" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dao" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbcore" - "github.com/chroma/chroma-coordinator/internal/notification" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" - "github.com/chroma/chroma-coordinator/internal/utils" + "github.com/chroma-core/chroma/go/pkg/coordinator" + "github.com/chroma-core/chroma/go/pkg/memberlist_manager" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/notification" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/utils" "github.com/pingcap/log" "go.uber.org/zap" "google.golang.org/grpc" @@ -29,13 +30,7 @@ type Config struct { SystemCatalogProvider string // MetaTable config - Username string - Password string - Address string - Port int - DBName string - MaxIdleConns int - MaxOpenConns int + DBConfig dbcore.DBConfig // Notification config NotificationStoreProvider string @@ -49,15 +44,19 @@ type Config struct { PulsarNamespace string // Kubernetes config - KubernetesNamespace string - WorkerMemberlistName string + KubernetesNamespace string - // Assignment policy config can be "simple" or "rendezvous" - AssignmentPolicy string + // Query service memberlist config + QueryServiceMemberlistName string + QueryServicePodLabel string // Watcher config WatchInterval time.Duration + // Compaction service memberlist config + CompactionServiceMemberlistName string + CompactionServicePodLabel string + // Config for testing Testing bool } @@ -77,16 +76,8 @@ func New(config Config) (*Server, error) { if config.SystemCatalogProvider == "memory" { return NewWithGrpcProvider(config, grpcutils.Default, nil) } else if config.SystemCatalogProvider == "database" { - dBConfig := dbcore.DBConfig{ - Username: config.Username, - Password: config.Password, - Address: config.Address, - Port: config.Port, - DBName: config.DBName, - MaxIdleConns: config.MaxIdleConns, - MaxOpenConns: config.MaxOpenConns, - } - db, err := dbcore.Connect(dBConfig) + dBConfig := config.DBConfig + db, err := dbcore.ConnectPostgres(dBConfig) if err != nil { return nil, err } @@ -102,22 +93,6 @@ func NewWithGrpcProvider(config Config, provider grpcutils.GrpcProvider, db *gor healthServer: health.NewServer(), } - var assignmentPolicy coordinator.CollectionAssignmentPolicy - if config.AssignmentPolicy == "simple" { - log.Info("Using simple assignment policy") - assignmentPolicy = coordinator.NewSimpleAssignmentPolicy(config.PulsarTenant, config.PulsarNamespace) - } else if config.AssignmentPolicy == "rendezvous" { - log.Info("Using rendezvous assignment policy") - err := utils.CreateTopics(config.PulsarAdminURL, config.PulsarTenant, config.PulsarNamespace, coordinator.Topics[:]) - if err != nil { - log.Error("Failed to create topics", zap.Error(err)) - return nil, err - } - assignmentPolicy = coordinator.NewRendezvousAssignmentPolicy(config.PulsarTenant, config.PulsarNamespace) - } else { - return nil, errors.New("invalid assignment policy, only simple and rendezvous are supported") - } - var notificationStore notification.NotificationStore if config.NotificationStoreProvider == "memory" { log.Info("Using memory notification store") @@ -157,25 +132,38 @@ func NewWithGrpcProvider(config Config, provider grpcutils.GrpcProvider, db *gor defer producer.Close() } - coordinator, err := coordinator.NewCoordinator(ctx, assignmentPolicy, db, notificationStore, notifier) + coordinator, err := coordinator.NewCoordinator(ctx, db, notificationStore, notifier) if err != nil { return nil, err } s.coordinator = coordinator s.coordinator.Start() if !config.Testing { - memberlist_manager, err := createMemberlistManager(config) + namespace := config.KubernetesNamespace + // Create memberlist manager for query service + queryMemberlistManager, err := createMemberlistManager(namespace, config.QueryServiceMemberlistName, config.QueryServicePodLabel, config.WatchInterval) + if err != nil { + return nil, err + } + + // Create memberlist manager for compaction service + compactionMemberlistManager, err := createMemberlistManager(namespace, config.CompactionServiceMemberlistName, config.CompactionServicePodLabel, config.WatchInterval) if err != nil { return nil, err } - // Start the memberlist manager - err = memberlist_manager.Start() + // Start the memberlist manager for query service + err = queryMemberlistManager.Start() + if err != nil { + return nil, err + } + // Start the memberlist manager for compaction service + err = compactionMemberlistManager.Start() if err != nil { return nil, err } - s.grpcServer, err = provider.StartGrpcServer("coordinator", config.GrpcConfig, func(registrar grpc.ServiceRegistrar) { + s.grpcServer, err = provider.StartGrpcServer("coordinator", config.GrpcConfig, func(registrar grpc.ServiceRegistrar) { coordinatorpb.RegisterSysDBServer(registrar, s) }) if err != nil { @@ -185,11 +173,8 @@ func NewWithGrpcProvider(config Config, provider grpcutils.GrpcProvider, db *gor return s, nil } -func createMemberlistManager(config Config) (*memberlist_manager.MemberlistManager, error) { - // TODO: Make this configuration - log.Info("Starting memberlist manager") - memberlist_name := config.WorkerMemberlistName - namespace := config.KubernetesNamespace +func createMemberlistManager(namespace string, memberlistName string, podLabel string, watchInterval time.Duration) (*memberlist_manager.MemberlistManager, error) { + log.Info("Creating memberlist manager for {}", zap.String("memberlist", memberlistName)) clientset, err := utils.GetKubernetesInterface() if err != nil { return nil, err @@ -198,8 +183,8 @@ func createMemberlistManager(config Config) (*memberlist_manager.MemberlistManag if err != nil { return nil, err } - nodeWatcher := memberlist_manager.NewKubernetesWatcher(clientset, namespace, "worker", config.WatchInterval) - memberlistStore := memberlist_manager.NewCRMemberlistStore(dynamicClient, namespace, memberlist_name) + nodeWatcher := memberlist_manager.NewKubernetesWatcher(clientset, namespace, podLabel, watchInterval) + memberlistStore := memberlist_manager.NewCRMemberlistStore(dynamicClient, namespace, memberlistName) memberlist_manager := memberlist_manager.NewMemberlistManager(nodeWatcher, memberlistStore) return memberlist_manager, nil } diff --git a/go/coordinator/internal/grpccoordinator/tenant_database_service.go b/go/pkg/coordinator/grpc/tenant_database_service.go similarity index 55% rename from go/coordinator/internal/grpccoordinator/tenant_database_service.go rename to go/pkg/coordinator/grpc/tenant_database_service.go index eb36b3de949..7ae4445fb08 100644 --- a/go/coordinator/internal/grpccoordinator/tenant_database_service.go +++ b/go/pkg/coordinator/grpc/tenant_database_service.go @@ -1,15 +1,19 @@ -package grpccoordinator +package grpc import ( "context" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/pingcap/log" + "go.uber.org/zap" + "google.golang.org/protobuf/types/known/emptypb" - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" ) -func (s *Server) CreateDatabase(ctx context.Context, req *coordinatorpb.CreateDatabaseRequest) (*coordinatorpb.ChromaResponse, error) { - res := &coordinatorpb.ChromaResponse{} +func (s *Server) CreateDatabase(ctx context.Context, req *coordinatorpb.CreateDatabaseRequest) (*coordinatorpb.CreateDatabaseResponse, error) { + res := &coordinatorpb.CreateDatabaseResponse{} createDatabase := &model.CreateDatabase{ ID: req.GetId(), Name: req.GetName(), @@ -51,8 +55,8 @@ func (s *Server) GetDatabase(ctx context.Context, req *coordinatorpb.GetDatabase return res, nil } -func (s *Server) CreateTenant(ctx context.Context, req *coordinatorpb.CreateTenantRequest) (*coordinatorpb.ChromaResponse, error) { - res := &coordinatorpb.ChromaResponse{} +func (s *Server) CreateTenant(ctx context.Context, req *coordinatorpb.CreateTenantRequest) (*coordinatorpb.CreateTenantResponse, error) { + res := &coordinatorpb.CreateTenantResponse{} createTenant := &model.CreateTenant{ Name: req.GetName(), } @@ -89,3 +93,29 @@ func (s *Server) GetTenant(ctx context.Context, req *coordinatorpb.GetTenantRequ res.Status = setResponseStatus(successCode) return res, nil } + +func (s *Server) SetLastCompactionTimeForTenant(ctx context.Context, req *coordinatorpb.SetLastCompactionTimeForTenantRequest) (*emptypb.Empty, error) { + err := s.coordinator.SetTenantLastCompactionTime(ctx, req.TenantLastCompactionTime.TenantId, req.TenantLastCompactionTime.LastCompactionTime) + if err != nil { + log.Error("error SetTenantLastCompactionTime", zap.Any("request", req.TenantLastCompactionTime), zap.Error(err)) + return nil, grpcutils.BuildInternalGrpcError(err.Error()) + } + return &emptypb.Empty{}, nil +} + +func (s *Server) GetLastCompactionTimeForTenant(ctx context.Context, req *coordinatorpb.GetLastCompactionTimeForTenantRequest) (*coordinatorpb.GetLastCompactionTimeForTenantResponse, error) { + res := &coordinatorpb.GetLastCompactionTimeForTenantResponse{} + tenantIDs := req.TenantId + tenants, err := s.coordinator.GetTenantsLastCompactionTime(ctx, tenantIDs) + if err != nil { + log.Error("error GetLastCompactionTimeForTenant", zap.Any("tenantIDs", tenantIDs), zap.Error(err)) + return nil, grpcutils.BuildInternalGrpcError(err.Error()) + } + for _, tenant := range tenants { + res.TenantLastCompactionTime = append(res.TenantLastCompactionTime, &coordinatorpb.TenantLastCompactionTime{ + TenantId: tenant.ID, + LastCompactionTime: tenant.LastCompactionTime, + }) + } + return res, nil +} diff --git a/go/pkg/coordinator/grpc/tenant_database_service_test.go b/go/pkg/coordinator/grpc/tenant_database_service_test.go new file mode 100644 index 00000000000..825765e2610 --- /dev/null +++ b/go/pkg/coordinator/grpc/tenant_database_service_test.go @@ -0,0 +1,107 @@ +package grpc + +import ( + "context" + "testing" + "time" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/chroma-core/chroma/go/pkg/metastore/coordinator" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "google.golang.org/genproto/googleapis/rpc/code" + codes "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "gorm.io/gorm" +) + +type TenantDatabaseServiceTestSuite struct { + suite.Suite + catalog *coordinator.Catalog + db *gorm.DB + s *Server +} + +func (suite *TenantDatabaseServiceTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() + s, err := NewWithGrpcProvider(Config{ + SystemCatalogProvider: "database", + NotificationStoreProvider: "memory", + NotifierProvider: "memory", + Testing: true}, grpcutils.Default, suite.db) + if err != nil { + suite.T().Fatalf("error creating server: %v", err) + } + suite.s = s + txnImpl := dbcore.NewTxImpl() + metaDomain := dao.NewMetaDomain() + suite.catalog = coordinator.NewTableCatalogWithNotification(txnImpl, metaDomain, nil) +} + +func (suite *TenantDatabaseServiceTestSuite) SetupTest() { + log.Info("setup test") +} + +func (suite *TenantDatabaseServiceTestSuite) TearDownTest() { + log.Info("teardown test") +} + +func (suite *TenantDatabaseServiceTestSuite) TestServer_TenantLastCompactionTime() { + log.Info("TestServer_TenantLastCompactionTime") + tenantId := "TestTenantLastCompactionTime" + request := &coordinatorpb.SetLastCompactionTimeForTenantRequest{ + TenantLastCompactionTime: &coordinatorpb.TenantLastCompactionTime{ + TenantId: tenantId, + LastCompactionTime: 0, + }, + } + _, err := suite.s.SetLastCompactionTimeForTenant(context.Background(), request) + suite.Equal(status.Error(codes.Code(code.Code_INTERNAL), common.ErrTenantNotFound.Error()), err) + + // create tenant + _, err = suite.catalog.CreateTenant(context.Background(), &model.CreateTenant{ + Name: tenantId, + Ts: time.Now().Unix(), + }, time.Now().Unix()) + if err != nil { + return + } + suite.NoError(err) + + _, err = suite.s.SetLastCompactionTimeForTenant(context.Background(), request) + suite.NoError(err) + tenants, err := suite.s.GetLastCompactionTimeForTenant(context.Background(), &coordinatorpb.GetLastCompactionTimeForTenantRequest{ + TenantId: []string{tenantId}, + }) + suite.NoError(err) + suite.Equal(1, len(tenants.TenantLastCompactionTime)) + suite.Equal(tenantId, tenants.TenantLastCompactionTime[0].TenantId) + suite.Equal(int64(0), tenants.TenantLastCompactionTime[0].LastCompactionTime) + + // update last compaction time + request.TenantLastCompactionTime.LastCompactionTime = 1 + _, err = suite.s.SetLastCompactionTimeForTenant(context.Background(), request) + suite.NoError(err) + tenants, err = suite.s.GetLastCompactionTimeForTenant(context.Background(), &coordinatorpb.GetLastCompactionTimeForTenantRequest{ + TenantId: []string{tenantId}, + }) + suite.NoError(err) + suite.Equal(1, len(tenants.TenantLastCompactionTime)) + suite.Equal(tenantId, tenants.TenantLastCompactionTime[0].TenantId) + suite.Equal(int64(1), tenants.TenantLastCompactionTime[0].LastCompactionTime) + + // clean up + err = dao.CleanUpTestTenant(suite.db, tenantId) + suite.NoError(err) +} + +func TestTenantDatabaseServiceTestSuite(t *testing.T) { + testSuite := new(TenantDatabaseServiceTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/coordinator/internal/grpccoordinator/grpcutils/config.go b/go/pkg/grpcutils/config.go similarity index 100% rename from go/coordinator/internal/grpccoordinator/grpcutils/config.go rename to go/pkg/grpcutils/config.go diff --git a/go/coordinator/internal/grpccoordinator/grpcutils/config_test.go b/go/pkg/grpcutils/config_test.go similarity index 100% rename from go/coordinator/internal/grpccoordinator/grpcutils/config_test.go rename to go/pkg/grpcutils/config_test.go diff --git a/go/pkg/grpcutils/response.go b/go/pkg/grpcutils/response.go new file mode 100644 index 00000000000..5a89344eb30 --- /dev/null +++ b/go/pkg/grpcutils/response.go @@ -0,0 +1,45 @@ +package grpcutils + +import ( + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/pingcap/log" + "go.uber.org/zap" + "google.golang.org/genproto/googleapis/rpc/errdetails" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func BuildInvalidArgumentGrpcError(fieldName string, desc string) (error, error) { + log.Info("InvalidArgument", zap.String("fieldName", fieldName), zap.String("desc", desc)) + st := status.New(codes.InvalidArgument, "invalid "+fieldName) + v := &errdetails.BadRequest_FieldViolation{ + Field: fieldName, + Description: desc, + } + br := &errdetails.BadRequest{ + FieldViolations: []*errdetails.BadRequest_FieldViolation{v}, + } + st, err := st.WithDetails(br) + if err != nil { + log.Error("Unexpected error attaching metadata", zap.Error(err)) + return nil, err + } + return st.Err(), nil +} + +func BuildInternalGrpcError(msg string) error { + return status.Error(codes.Internal, msg) +} + +func BuildErrorForUUID(ID types.UniqueID, name string, err error) error { + if err != nil || ID == types.NilUniqueID() { + log.Error(name+"id format error", zap.String(name+".id", ID.String())) + grpcError, err := BuildInvalidArgumentGrpcError(name+"_id", "wrong "+name+"_id format") + if err != nil { + log.Error("error building grpc error", zap.Error(err)) + return err + } + return grpcError + } + return nil +} diff --git a/go/coordinator/internal/grpccoordinator/grpcutils/service.go b/go/pkg/grpcutils/service.go similarity index 94% rename from go/coordinator/internal/grpccoordinator/grpcutils/service.go rename to go/pkg/grpcutils/service.go index e721f2158f9..885726e34c5 100644 --- a/go/coordinator/internal/grpccoordinator/grpcutils/service.go +++ b/go/pkg/grpcutils/service.go @@ -3,6 +3,7 @@ package grpcutils import ( "crypto/tls" "crypto/x509" + "github.com/chroma-core/chroma/go/shared/otel" "io" "net" "os" @@ -70,6 +71,7 @@ func newDefaultGrpcProvider(name string, grpcConfig *GrpcConfig, registerFunc fu opts = append(opts, grpc.Creds(credentials.NewTLS(tlsConfig))) } + opts = append(opts, grpc.UnaryInterceptor(otel.ServerGrpcInterceptor)) c := &defaultGrpcServer{ server: grpc.NewServer(opts...), diff --git a/go/pkg/logservice/apis.go b/go/pkg/logservice/apis.go new file mode 100644 index 00000000000..40736ee6cfa --- /dev/null +++ b/go/pkg/logservice/apis.go @@ -0,0 +1,32 @@ +package logservice + +import ( + "context" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" +) + +type ( + IRecordLog interface { + common.Component + PushLogs(ctx context.Context, collectionID types.UniqueID, recordContent [][]byte) (int, error) + PullLogs(ctx context.Context, collectionID types.UniqueID, id int64, batchSize int, endTimestamp int64) ([]*dbmodel.RecordLog, error) + GetAllCollectionIDsToCompact() ([]*dbmodel.RecordLog, error) + } +) + +var _ IRecordLog = &RecordLog{} + +func (s *RecordLog) PushLogs(ctx context.Context, collectionID types.UniqueID, recordsContent [][]byte) (int, error) { + return s.recordLogDb.PushLogs(collectionID, recordsContent) +} + +func (s *RecordLog) PullLogs(ctx context.Context, collectionID types.UniqueID, id int64, batchSize int, endTimestamp int64) ([]*dbmodel.RecordLog, error) { + return s.recordLogDb.PullLogs(collectionID, id, batchSize, endTimestamp) +} + +func (s *RecordLog) GetAllCollectionIDsToCompact() ([]*dbmodel.RecordLog, error) { + return s.recordLogDb.GetAllCollectionsToCompact() +} diff --git a/go/pkg/logservice/grpc/record_log_service.go b/go/pkg/logservice/grpc/record_log_service.go new file mode 100644 index 00000000000..39461110624 --- /dev/null +++ b/go/pkg/logservice/grpc/record_log_service.go @@ -0,0 +1,106 @@ +package grpc + +import ( + "context" + + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/proto/logservicepb" + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/pingcap/log" + "go.uber.org/zap" + "google.golang.org/protobuf/proto" +) + +type CollectionInfo struct { + CollectionId string + FirstLogId int64 + FirstLogTs int64 +} + +func (s *Server) PushLogs(ctx context.Context, req *logservicepb.PushLogsRequest) (*logservicepb.PushLogsResponse, error) { + res := &logservicepb.PushLogsResponse{} + collectionID, err := types.ToUniqueID(&req.CollectionId) + err = grpcutils.BuildErrorForUUID(collectionID, "collection", err) + if err != nil { + return nil, err + } + var recordsContent [][]byte + for _, record := range req.Records { + data, err := proto.Marshal(record) + if err != nil { + log.Error("marshaling error", zap.Error(err)) + grpcError, err := grpcutils.BuildInvalidArgumentGrpcError("records", "marshaling error") + if err != nil { + return nil, err + } + return nil, grpcError + } + recordsContent = append(recordsContent, data) + } + recordCount, err := s.logService.PushLogs(ctx, collectionID, recordsContent) + if err != nil { + log.Error("error pushing logs", zap.Error(err)) + return nil, grpcutils.BuildInternalGrpcError(err.Error()) + } + res.RecordCount = int32(recordCount) + log.Info("PushLogs success", zap.String("collectionID", req.CollectionId), zap.Int("recordCount", recordCount)) + return res, nil +} + +func (s *Server) PullLogs(ctx context.Context, req *logservicepb.PullLogsRequest) (*logservicepb.PullLogsResponse, error) { + res := &logservicepb.PullLogsResponse{} + collectionID, err := types.ToUniqueID(&req.CollectionId) + err = grpcutils.BuildErrorForUUID(collectionID, "collection", err) + if err != nil { + return nil, err + } + records := make([]*logservicepb.LogRecord, 0) + recordLogs, err := s.logService.PullLogs(ctx, collectionID, req.GetStartFromOffset(), int(req.BatchSize), req.GetEndTimestamp()) + if err != nil { + log.Error("error pulling logs", zap.Error(err)) + return nil, grpcutils.BuildInternalGrpcError(err.Error()) + } + for index := range recordLogs { + record := &coordinatorpb.OperationRecord{} + if err := proto.Unmarshal(*recordLogs[index].Record, record); err != nil { + log.Error("Unmarshal error", zap.Error(err)) + grpcError, err := grpcutils.BuildInvalidArgumentGrpcError("records", "marshaling error") + if err != nil { + return nil, err + } + return nil, grpcError + } + recordLog := &logservicepb.LogRecord{ + LogOffset: recordLogs[index].LogOffset, + Record: record, + } + records = append(records, recordLog) + } + res.Records = records + log.Info("PullLogs success", zap.String("collectionID", req.CollectionId), zap.Int("recordCount", len(records))) + return res, nil +} + +func (s *Server) GetAllCollectionInfoToCompact(ctx context.Context, req *logservicepb.GetAllCollectionInfoToCompactRequest) (*logservicepb.GetAllCollectionInfoToCompactResponse, error) { + res := &logservicepb.GetAllCollectionInfoToCompactResponse{} + res.AllCollectionInfo = make([]*logservicepb.CollectionInfo, 0) + var recordLogs []*dbmodel.RecordLog + recordLogs, err := s.logService.GetAllCollectionIDsToCompact() + if err != nil { + log.Error("error getting collection info", zap.Error(err)) + return nil, grpcutils.BuildInternalGrpcError(err.Error()) + } + for _, recordLog := range recordLogs { + collectionInfo := &logservicepb.CollectionInfo{ + CollectionId: *recordLog.CollectionID, + FirstLogOffset: recordLog.LogOffset, + FirstLogTs: recordLog.Timestamp, + } + res.AllCollectionInfo = append(res.AllCollectionInfo, collectionInfo) + } + // print everything for now, we can make this smaller once + log.Info("GetAllCollectionInfoToCompact success", zap.Any("collectionInfo", res.AllCollectionInfo)) + return res, nil +} diff --git a/go/pkg/logservice/grpc/record_log_service_test.go b/go/pkg/logservice/grpc/record_log_service_test.go new file mode 100644 index 00000000000..a71d62976e6 --- /dev/null +++ b/go/pkg/logservice/grpc/record_log_service_test.go @@ -0,0 +1,231 @@ +package grpc + +import ( + "bytes" + "context" + "encoding/binary" + "testing" + "time" + + "github.com/chroma-core/chroma/go/pkg/logservice/testutils" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/proto/logservicepb" + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + "gorm.io/gorm" +) + +type RecordLogServiceTestSuite struct { + suite.Suite + db *gorm.DB + s *Server + collectionId types.UniqueID +} + +func (suite *RecordLogServiceTestSuite) SetupSuite() { + log.Info("setup suite") + // setup server and db + s, _ := New(Config{ + DBProvider: "postgres", + DBConfig: dbcore.GetDBConfigForTesting(), + StartGrpc: false, + }) + suite.s = s + suite.db = dbcore.ConfigDatabaseForTesting() + recordLogTableExist := suite.db.Migrator().HasTable(&dbmodel.RecordLog{}) + if !recordLogTableExist { + err := suite.db.Migrator().CreateTable(&dbmodel.RecordLog{}) + suite.NoError(err) + } +} + +func (suite *RecordLogServiceTestSuite) SetupTest() { + log.Info("setup test") + suite.collectionId = types.NewUniqueID() + err := testutils.CreateCollections(suite.db, suite.collectionId) + suite.NoError(err) +} + +func (suite *RecordLogServiceTestSuite) TearDownTest() { + log.Info("teardown test") + err := testutils.CleanupCollections(suite.db, suite.collectionId) + suite.NoError(err) +} + +func encodeVector(dimension int32, vector []float32, encoding coordinatorpb.ScalarEncoding) *coordinatorpb.Vector { + buf := new(bytes.Buffer) + err := binary.Write(buf, binary.LittleEndian, vector) + if err != nil { + panic(err) + } + + return &coordinatorpb.Vector{ + Dimension: dimension, + Vector: buf.Bytes(), + Encoding: encoding, + } +} + +func GetTestEmbeddingRecords(collectionId string) (recordsToSubmit []*coordinatorpb.OperationRecord) { + testVector1 := []float32{1.0, 2.0, 3.0} + testVector2 := []float32{1.2, 2.24, 3.2} + testVector3 := []float32{7.0, 8.0, 9.0} + recordsToSubmit = make([]*coordinatorpb.OperationRecord, 0) + recordsToSubmit = append(recordsToSubmit, &coordinatorpb.OperationRecord{ + Id: types.NewUniqueID().String(), + Vector: encodeVector(10, testVector1, coordinatorpb.ScalarEncoding_FLOAT32), + Operation: coordinatorpb.Operation_ADD, + }) + recordsToSubmit = append(recordsToSubmit, &coordinatorpb.OperationRecord{ + Id: types.NewUniqueID().String(), + Vector: encodeVector(6, testVector2, coordinatorpb.ScalarEncoding_FLOAT32), + Operation: coordinatorpb.Operation_UPDATE, + }) + recordsToSubmit = append(recordsToSubmit, &coordinatorpb.OperationRecord{ + Id: types.NewUniqueID().String(), + Vector: encodeVector(10, testVector3, coordinatorpb.ScalarEncoding_FLOAT32), + Operation: coordinatorpb.Operation_ADD, + }) + return recordsToSubmit +} + +func (suite *RecordLogServiceTestSuite) TestServer_PushLogs() { + log.Info("test push logs") + // push some records + recordsToSubmit := GetTestEmbeddingRecords(suite.collectionId.String()) + pushRequest := logservicepb.PushLogsRequest{ + CollectionId: suite.collectionId.String(), + Records: recordsToSubmit, + } + response, err := suite.s.PushLogs(context.Background(), &pushRequest) + suite.NoError(err) + suite.Equal(int32(3), response.RecordCount) + + var recordLogs []*dbmodel.RecordLog + suite.db.Where("collection_id = ?", types.FromUniqueID(suite.collectionId)).Find(&recordLogs) + suite.Len(recordLogs, 3) + for index := range recordLogs { + suite.Equal(int64(index+1), recordLogs[index].LogOffset) + suite.Equal(suite.collectionId.String(), *recordLogs[index].CollectionID) + record := &coordinatorpb.OperationRecord{} + if unmarshalErr := proto.Unmarshal(*recordLogs[index].Record, record); err != nil { + suite.NoError(unmarshalErr) + } + suite.Equal(recordsToSubmit[index].Id, record.Id) + suite.Equal(recordsToSubmit[index].Operation, record.Operation) + suite.Equal(recordsToSubmit[index].Metadata, record.Metadata) + suite.Equal(recordsToSubmit[index].Vector.Dimension, record.Vector.Dimension) + suite.Equal(recordsToSubmit[index].Vector.Encoding, record.Vector.Encoding) + suite.Equal(recordsToSubmit[index].Vector.Vector, record.Vector.Vector) + } +} + +func (suite *RecordLogServiceTestSuite) TestServer_PullLogs() { + // push some records + recordsToSubmit := GetTestEmbeddingRecords(suite.collectionId.String()) + // deep clone the records since PushLogs will mutate the records and we need a source of truth + recordsToSubmit_sot := make([]*coordinatorpb.OperationRecord, len(recordsToSubmit)) + for i := range recordsToSubmit { + recordsToSubmit_sot[i] = proto.Clone(recordsToSubmit[i]).(*coordinatorpb.OperationRecord) + } + pushRequest := logservicepb.PushLogsRequest{ + CollectionId: suite.collectionId.String(), + Records: recordsToSubmit, + } + _, err := suite.s.PushLogs(context.Background(), &pushRequest) + suite.NoError(err) + + // pull the records + pullRequest := logservicepb.PullLogsRequest{ + CollectionId: suite.collectionId.String(), + StartFromOffset: 0, + BatchSize: 10, + } + pullResponse, err := suite.s.PullLogs(context.Background(), &pullRequest) + suite.NoError(err) + suite.Len(pullResponse.Records, 3) + for index := range pullResponse.Records { + suite.Equal(int64(index+1), pullResponse.Records[index].LogOffset) + suite.Equal(recordsToSubmit_sot[index].Id, pullResponse.Records[index].Record.Id) + suite.Equal(recordsToSubmit_sot[index].Operation, pullResponse.Records[index].Record.Operation) + suite.Equal(recordsToSubmit_sot[index].Metadata, pullResponse.Records[index].Record.Metadata) + suite.Equal(recordsToSubmit_sot[index].Vector.Dimension, pullResponse.Records[index].Record.Vector.Dimension) + suite.Equal(recordsToSubmit_sot[index].Vector.Encoding, pullResponse.Records[index].Record.Vector.Encoding) + suite.Equal(recordsToSubmit_sot[index].Vector.Vector, pullResponse.Records[index].Record.Vector.Vector) + } +} + +func (suite *RecordLogServiceTestSuite) TestServer_Bad_CollectionId() { + log.Info("test bad collectionId") + // push some records + pushRequest := logservicepb.PushLogsRequest{ + CollectionId: "badId", + Records: []*coordinatorpb.OperationRecord{}, + } + _, err := suite.s.PushLogs(context.Background(), &pushRequest) + suite.Error(err) + st, ok := status.FromError(err) + suite.True(ok) + suite.Equal(codes.InvalidArgument, st.Code()) + suite.Equal("invalid collection_id", st.Message()) + + // pull the records + // pull the records + pullRequest := logservicepb.PullLogsRequest{ + CollectionId: "badId", + StartFromOffset: 0, + BatchSize: 10, + } + _, err = suite.s.PullLogs(context.Background(), &pullRequest) + suite.Error(err) + st, ok = status.FromError(err) + suite.True(ok) + suite.Equal(codes.InvalidArgument, st.Code()) + suite.Equal("invalid collection_id", st.Message()) +} + +func (suite *RecordLogServiceTestSuite) TestServer_GetAllCollectionInfoToCompact() { + // push some records + var startTime = time.Now().UnixNano() + recordsToSubmit := GetTestEmbeddingRecords(suite.collectionId.String()) + pushRequest := logservicepb.PushLogsRequest{ + CollectionId: suite.collectionId.String(), + Records: recordsToSubmit, + } + _, err := suite.s.PushLogs(context.Background(), &pushRequest) + suite.NoError(err) + + // get collection info for compactor + request := logservicepb.GetAllCollectionInfoToCompactRequest{} + response, err := suite.s.GetAllCollectionInfoToCompact(context.Background(), &request) + suite.NoError(err) + suite.Len(response.AllCollectionInfo, 1) + suite.Equal(suite.collectionId.String(), response.AllCollectionInfo[0].CollectionId) + suite.Equal(int64(1), response.AllCollectionInfo[0].FirstLogOffset) + suite.True(response.AllCollectionInfo[0].FirstLogTs > startTime) + suite.True(response.AllCollectionInfo[0].FirstLogTs < time.Now().UnixNano()) + + // move log position + testutils.MoveLogPosition(suite.db, suite.collectionId, 2) + + // get collection info for compactor + request = logservicepb.GetAllCollectionInfoToCompactRequest{} + response, err = suite.s.GetAllCollectionInfoToCompact(context.Background(), &request) + suite.NoError(err) + suite.Len(response.AllCollectionInfo, 1) + suite.Equal(suite.collectionId.String(), response.AllCollectionInfo[0].CollectionId) + suite.Equal(int64(3), response.AllCollectionInfo[0].FirstLogOffset) + suite.True(response.AllCollectionInfo[0].FirstLogTs > startTime) + suite.True(response.AllCollectionInfo[0].FirstLogTs < time.Now().UnixNano()) +} + +func TestRecordLogServiceTestSuite(t *testing.T) { + testSuite := new(RecordLogServiceTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/pkg/logservice/grpc/server.go b/go/pkg/logservice/grpc/server.go new file mode 100644 index 00000000000..c38c12a7bb0 --- /dev/null +++ b/go/pkg/logservice/grpc/server.go @@ -0,0 +1,104 @@ +package grpc + +import ( + "context" + "errors" + "github.com/chroma-core/chroma/go/pkg/grpcutils" + "github.com/chroma-core/chroma/go/pkg/logservice" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/proto/logservicepb" + "github.com/pingcap/log" + "go.uber.org/zap" + "google.golang.org/grpc" + "google.golang.org/grpc/health" +) + +type Config struct { + // GrpcConfig config + GrpcConfig *grpcutils.GrpcConfig + + // System catalog provider + DBProvider string + + // Postgres config + DBConfig dbcore.DBConfig + + // whether to start grpc service + StartGrpc bool +} + +type Server struct { + logservicepb.UnimplementedLogServiceServer + logService logservice.IRecordLog + grpcServer grpcutils.GrpcServer + healthServer *health.Server +} + +func New(config Config) (*Server, error) { + log.Info("New Log Service...") + + if config.DBProvider == "postgres" { + dBConfig := config.DBConfig + _, err := dbcore.ConnectPostgres(dBConfig) + if err != nil { + log.Error("Error connecting to Postgres DB.", zap.Error(err)) + panic(err) + } + } else { + log.Error("invalid DB provider, only postgres is supported") + return nil, errors.New("invalid DB provider, only postgres is supported") + } + + s := startLogService() + if config.StartGrpc { + s.grpcServer = startGrpcService(s, config.GrpcConfig) + } + + log.Info("New Log Service Completed.") + return s, nil +} + +func startLogService() *Server { + log.Info("Staring Log Service...") + ctx := context.Background() + s := &Server{ + healthServer: health.NewServer(), + } + + logService, err := logservice.NewLogService(ctx) + if err != nil { + log.Error("Error creating Log Service.", zap.Error(err)) + panic(err) + } + s.logService = logService + err = s.logService.Start() + if err != nil { + log.Error("Error starting Log Service.", zap.Error(err)) + panic(err) + } + log.Info("Log Service Started.") + return s +} + +func startGrpcService(s *Server, grpcConfig *grpcutils.GrpcConfig) grpcutils.GrpcServer { + log.Info("Staring Grpc Service...") + server, err := grpcutils.Default.StartGrpcServer("logservice", grpcConfig, func(registrar grpc.ServiceRegistrar) { + logservicepb.RegisterLogServiceServer(registrar, s) + }) + if err != nil { + log.Error("Error starting grpc Service.", zap.Error(err)) + panic(err) + } + return server +} + +func (s *Server) Close() error { + s.healthServer.Shutdown() + err := s.logService.Stop() + if err != nil { + log.Error("Failed to stop log service", zap.Error(err)) + return err + } + log.Info("Server closed") + return nil +} diff --git a/go/pkg/logservice/recordlog.go b/go/pkg/logservice/recordlog.go new file mode 100644 index 00000000000..5207b4f81fa --- /dev/null +++ b/go/pkg/logservice/recordlog.go @@ -0,0 +1,33 @@ +package logservice + +import ( + "context" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/pingcap/log" +) + +var _ IRecordLog = (*RecordLog)(nil) + +type RecordLog struct { + ctx context.Context + recordLogDb dbmodel.IRecordLogDb +} + +func NewLogService(ctx context.Context) (*RecordLog, error) { + s := &RecordLog{ + ctx: ctx, + recordLogDb: dao.NewMetaDomain().RecordLogDb(ctx), + } + return s, nil +} + +func (s *RecordLog) Start() error { + log.Info("RecordLog start") + return nil +} + +func (s *RecordLog) Stop() error { + log.Info("RecordLog stop") + return nil +} diff --git a/go/pkg/logservice/testutils/record_log_test_util.go b/go/pkg/logservice/testutils/record_log_test_util.go new file mode 100644 index 00000000000..9f543509889 --- /dev/null +++ b/go/pkg/logservice/testutils/record_log_test_util.go @@ -0,0 +1,49 @@ +package testutils + +import ( + "strconv" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" + "gorm.io/gorm" +) + +func CreateCollections(db *gorm.DB, collectionIds ...types.UniqueID) error { + // create test collections + for index, collectionId := range collectionIds { + collectionName := "collection" + strconv.Itoa(index+1) + var collectionDimension int32 = 6 + collection := &dbmodel.Collection{ + ID: collectionId.String(), + Name: &collectionName, + Dimension: &collectionDimension, + DatabaseID: types.NewUniqueID().String(), + } + err := db.Create(collection).Error + if err != nil { + return err + } + } + return nil +} + +func CleanupCollections(db *gorm.DB, collectionIds ...types.UniqueID) error { + // delete test collections + for _, collectionId := range collectionIds { + err := db.Where("id = ?", collectionId.String()).Delete(&dbmodel.Collection{}).Error + if err != nil { + return err + } + } + + // cleanup logs + err := db.Where("collection_id in ?", collectionIds).Delete(&dbmodel.RecordLog{}).Error + if err != nil { + return err + } + return nil +} + +func MoveLogPosition(db *gorm.DB, collectionId types.UniqueID, position int64) { + db.Model(&dbmodel.Collection{}).Where("id = ?", collectionId.String()).Update("log_position", position) +} diff --git a/go/coordinator/internal/memberlist_manager/memberlist_manager.go b/go/pkg/memberlist_manager/memberlist_manager.go similarity index 94% rename from go/coordinator/internal/memberlist_manager/memberlist_manager.go rename to go/pkg/memberlist_manager/memberlist_manager.go index 3da53fbc3b9..e1afd4dc70a 100644 --- a/go/coordinator/internal/memberlist_manager/memberlist_manager.go +++ b/go/pkg/memberlist_manager/memberlist_manager.go @@ -4,7 +4,7 @@ import ( "context" "errors" - "github.com/chroma/chroma-coordinator/internal/common" + "github.com/chroma-core/chroma/go/pkg/common" "github.com/pingcap/log" "go.uber.org/zap" "k8s.io/client-go/util/workqueue" @@ -39,6 +39,7 @@ func NewMemberlistManager(nodeWatcher IWatcher, memberlistStore IMemberlistStore } func (m *MemberlistManager) Start() error { + log.Info("Starting memberlist manager") m.nodeWatcher.RegisterCallback(func(nodeIp string) { m.workqueue.Add(nodeIp) }) @@ -110,7 +111,7 @@ func (m *MemberlistManager) reconcile(nodeIp string, status Status) error { if !exists && status == Ready { newMemberlist = append(newMemberlist, nodeIp) } - return m.memberlistStore.UpdateMemberlist(context.TODO(), &newMemberlist, resourceVersion) + return m.memberlistStore.UpdateMemberlist(context.Background(), &newMemberlist, resourceVersion) } func (m *MemberlistManager) Stop() error { diff --git a/go/coordinator/internal/memberlist_manager/memberlist_manager_test.go b/go/pkg/memberlist_manager/memberlist_manager_test.go similarity index 99% rename from go/coordinator/internal/memberlist_manager/memberlist_manager_test.go rename to go/pkg/memberlist_manager/memberlist_manager_test.go index 4a26fdd484b..9e6ad52119c 100644 --- a/go/coordinator/internal/memberlist_manager/memberlist_manager_test.go +++ b/go/pkg/memberlist_manager/memberlist_manager_test.go @@ -6,7 +6,7 @@ import ( "testing" "time" - "github.com/chroma/chroma-coordinator/internal/utils" + "github.com/chroma-core/chroma/go/pkg/utils" "github.com/stretchr/testify/assert" v1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" diff --git a/go/coordinator/internal/memberlist_manager/memberlist_store.go b/go/pkg/memberlist_manager/memberlist_store.go similarity index 96% rename from go/coordinator/internal/memberlist_manager/memberlist_store.go rename to go/pkg/memberlist_manager/memberlist_store.go index 0567897f46e..49191359d0c 100644 --- a/go/coordinator/internal/memberlist_manager/memberlist_store.go +++ b/go/pkg/memberlist_manager/memberlist_store.go @@ -3,6 +3,8 @@ package memberlist_manager import ( "context" + "github.com/pingcap/log" + "go.uber.org/zap" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/runtime/schema" @@ -53,6 +55,7 @@ func (s *CRMemberlistStore) GetMemberlist(ctx context.Context) (return_memberlis func (s *CRMemberlistStore) UpdateMemberlist(ctx context.Context, memberlist *Memberlist, resourceVersion string) error { gvr := getGvr() + log.Info("Updating memberlist store", zap.Any("memberlist", memberlist)) unstructured := memberlistToCr(memberlist, s.coordinatorNamespace, s.memberlistCustomResource, resourceVersion) _, err := s.dynamicClient.Resource(gvr).Namespace("chroma").Update(context.TODO(), unstructured, metav1.UpdateOptions{}) if err != nil { diff --git a/go/coordinator/internal/memberlist_manager/node_watcher.go b/go/pkg/memberlist_manager/node_watcher.go similarity index 90% rename from go/coordinator/internal/memberlist_manager/node_watcher.go rename to go/pkg/memberlist_manager/node_watcher.go index d534620eeb9..d3d2a04944b 100644 --- a/go/coordinator/internal/memberlist_manager/node_watcher.go +++ b/go/pkg/memberlist_manager/node_watcher.go @@ -5,7 +5,7 @@ import ( "sync" "time" - "github.com/chroma/chroma-coordinator/internal/common" + "github.com/chroma-core/chroma/go/pkg/common" "github.com/pingcap/log" "go.uber.org/zap" v1 "k8s.io/api/core/v1" @@ -47,6 +47,7 @@ type KubernetesWatcher struct { } func NewKubernetesWatcher(clientset kubernetes.Interface, coordinator_namespace string, pod_label string, resyncPeriod time.Duration) *KubernetesWatcher { + log.Info("Creating new kubernetes watcher", zap.String("namespace", coordinator_namespace), zap.String("pod label", pod_label), zap.Duration("resync period", resyncPeriod)) labelSelector := labels.SelectorFromSet(map[string]string{MemberLabel: pod_label}) factory := informers.NewSharedInformerFactoryWithOptions(clientset, resyncPeriod, informers.WithNamespace(coordinator_namespace), informers.WithTweakListOptions(func(options *metav1.ListOptions) { options.LabelSelector = labelSelector.String() })) podInformer := factory.Core().V1().Pods().Informer() @@ -75,6 +76,7 @@ func (w *KubernetesWatcher) Start() error { log.Error("Error while asserting object to pod") } if err == nil { + log.Info("Kubernetes Pod Added", zap.String("key", key), zap.String("ip", objPod.Status.PodIP)) ip := objPod.Status.PodIP w.mu.Lock() w.ipToKey[ip] = key @@ -91,6 +93,7 @@ func (w *KubernetesWatcher) Start() error { log.Error("Error while asserting object to pod") } if err == nil { + log.Info("Kubernetes Pod Updated", zap.String("key", key), zap.String("ip", objPod.Status.PodIP)) ip := objPod.Status.PodIP w.ipToKey[ip] = key w.notify(ip) @@ -105,6 +108,7 @@ func (w *KubernetesWatcher) Start() error { log.Error("Error while asserting object to pod") } if err == nil { + log.Info("Kubernetes Pod Deleted", zap.String("ip", objPod.Status.PodIP)) ip := objPod.Status.PodIP // The contract for GetStatus is that if the ip is not in this map, then it returns NotReady delete(w.ipToKey, ip) diff --git a/go/coordinator/internal/metastore/catalog.go b/go/pkg/metastore/catalog.go similarity index 70% rename from go/coordinator/internal/metastore/catalog.go rename to go/pkg/metastore/catalog.go index 8a54ebbf910..52a5d91037a 100644 --- a/go/coordinator/internal/metastore/catalog.go +++ b/go/pkg/metastore/catalog.go @@ -3,8 +3,10 @@ package metastore import ( "context" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/types" ) // Catalog defines methods for system catalog @@ -13,11 +15,11 @@ import ( type Catalog interface { ResetState(ctx context.Context) error CreateCollection(ctx context.Context, createCollection *model.CreateCollection, ts types.Timestamp) (*model.Collection, error) - GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*model.Collection, error) + GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, tenantID string, databaseName string) ([]*model.Collection, error) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection, ts types.Timestamp) (*model.Collection, error) CreateSegment(ctx context.Context, createSegment *model.CreateSegment, ts types.Timestamp) (*model.Segment, error) - GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID, ts types.Timestamp) ([]*model.Segment, error) + GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*model.Segment, error) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error UpdateSegment(ctx context.Context, segmentInfo *model.UpdateSegment, ts types.Timestamp) (*model.Segment, error) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase, ts types.Timestamp) (*model.Database, error) @@ -26,4 +28,7 @@ type Catalog interface { CreateTenant(ctx context.Context, createTenant *model.CreateTenant, ts types.Timestamp) (*model.Tenant, error) GetTenants(ctx context.Context, getTenant *model.GetTenant, ts types.Timestamp) (*model.Tenant, error) GetAllTenants(ctx context.Context, ts types.Timestamp) ([]*model.Tenant, error) + SetTenantLastCompactionTime(ctx context.Context, tenantID string, lastCompactionTime int64) error + GetTenantsLastCompactionTime(ctx context.Context, tenantIDs []string) ([]*dbmodel.Tenant, error) + FlushCollectionCompaction(ctx context.Context, flushCollectionCompaction *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error) } diff --git a/go/coordinator/internal/metastore/coordinator/model_db_convert.go b/go/pkg/metastore/coordinator/model_db_convert.go similarity index 95% rename from go/coordinator/internal/metastore/coordinator/model_db_convert.go rename to go/pkg/metastore/coordinator/model_db_convert.go index f5fb51bcaea..7c57feb896d 100644 --- a/go/coordinator/internal/metastore/coordinator/model_db_convert.go +++ b/go/pkg/metastore/coordinator/model_db_convert.go @@ -1,9 +1,9 @@ package coordinator import ( - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" ) @@ -17,11 +17,12 @@ func convertCollectionToModel(collectionAndMetadataList []*dbmodel.CollectionAnd collection := &model.Collection{ ID: types.MustParse(collectionAndMetadata.Collection.ID), Name: *collectionAndMetadata.Collection.Name, - Topic: *collectionAndMetadata.Collection.Topic, Dimension: collectionAndMetadata.Collection.Dimension, TenantID: collectionAndMetadata.TenantID, DatabaseName: collectionAndMetadata.DatabaseName, Ts: collectionAndMetadata.Collection.Ts, + LogPosition: collectionAndMetadata.Collection.LogPosition, + Version: collectionAndMetadata.Collection.Version, } collection.Metadata = convertCollectionMetadataToModel(collectionAndMetadata.CollectionMetadata) collections = append(collections, collection) @@ -96,7 +97,6 @@ func convertSegmentToModel(segmentAndMetadataList []*dbmodel.SegmentAndMetadata) ID: types.MustParse(segmentAndMetadata.Segment.ID), Type: segmentAndMetadata.Segment.Type, Scope: segmentAndMetadata.Segment.Scope, - Topic: segmentAndMetadata.Segment.Topic, Ts: segmentAndMetadata.Segment.Ts, } if segmentAndMetadata.Segment.CollectionID != nil { diff --git a/go/coordinator/internal/metastore/coordinator/model_db_convert_test.go b/go/pkg/metastore/coordinator/model_db_convert_test.go similarity index 93% rename from go/coordinator/internal/metastore/coordinator/model_db_convert_test.go rename to go/pkg/metastore/coordinator/model_db_convert_test.go index 67da68b1a76..27d4c1dd4da 100644 --- a/go/coordinator/internal/metastore/coordinator/model_db_convert_test.go +++ b/go/pkg/metastore/coordinator/model_db_convert_test.go @@ -4,9 +4,9 @@ import ( "sort" "testing" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/stretchr/testify/assert" ) @@ -76,14 +76,12 @@ func TestConvertSegmentToModel(t *testing.T) { // Test case 3: segmentAndMetadataList contains one segment with all fields set segmentID := types.MustParse("515fc331-e117-4b86-bd84-85341128c337") - segmentTopic := "segment_topic" collectionID := "d9a75e2e-2929-45c4-af06-75b15630edd0" segmentAndMetadata := &dbmodel.SegmentAndMetadata{ Segment: &dbmodel.Segment{ ID: segmentID.String(), Type: "segment_type", Scope: "segment_scope", - Topic: &segmentTopic, CollectionID: &collectionID, }, SegmentMetadata: []*dbmodel.SegmentMetadata{}, @@ -94,7 +92,6 @@ func TestConvertSegmentToModel(t *testing.T) { assert.Equal(t, segmentID, modelSegments[0].ID) assert.Equal(t, "segment_type", modelSegments[0].Type) assert.Equal(t, "segment_scope", modelSegments[0].Scope) - assert.Equal(t, "segment_topic", *modelSegments[0].Topic) assert.Equal(t, types.MustParse(collectionID), modelSegments[0].CollectionID) assert.Nil(t, modelSegments[0].Metadata) } @@ -136,13 +133,11 @@ func TestConvertCollectionToModel(t *testing.T) { // Test case 3: collectionAndMetadataList contains one collection with all fields set collectionID := types.MustParse("d9a75e2e-2929-45c4-af06-75b15630edd0") collectionName := "collection_name" - collectionTopic := "collection_topic" collectionDimension := int32(3) collectionAndMetadata := &dbmodel.CollectionAndMetadata{ Collection: &dbmodel.Collection{ ID: collectionID.String(), Name: &collectionName, - Topic: &collectionTopic, Dimension: &collectionDimension, }, CollectionMetadata: []*dbmodel.CollectionMetadata{}, @@ -152,7 +147,6 @@ func TestConvertCollectionToModel(t *testing.T) { assert.Len(t, modelCollections, 1) assert.Equal(t, collectionID, modelCollections[0].ID) assert.Equal(t, collectionName, modelCollections[0].Name) - assert.Equal(t, collectionTopic, modelCollections[0].Topic) assert.Equal(t, collectionDimension, *modelCollections[0].Dimension) assert.Nil(t, modelCollections[0].Metadata) } diff --git a/go/coordinator/internal/metastore/coordinator/table_catalog.go b/go/pkg/metastore/coordinator/table_catalog.go similarity index 73% rename from go/coordinator/internal/metastore/coordinator/table_catalog.go rename to go/pkg/metastore/coordinator/table_catalog.go index 4bd0d7f1244..fcbe8e76774 100644 --- a/go/coordinator/internal/metastore/coordinator/table_catalog.go +++ b/go/pkg/metastore/coordinator/table_catalog.go @@ -2,13 +2,14 @@ package coordinator import ( "context" - - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/notification" - "github.com/chroma/chroma-coordinator/internal/types" + "time" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/notification" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" ) @@ -37,19 +38,14 @@ var _ metastore.Catalog = (*Catalog)(nil) func (tc *Catalog) ResetState(ctx context.Context) error { return tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { - err := tc.metaDomain.CollectionDb(txCtx).DeleteAll() - if err != nil { - log.Error("error reset collection db", zap.Error(err)) - return err - } - err = tc.metaDomain.CollectionMetadataDb(txCtx).DeleteAll() + err := tc.metaDomain.CollectionMetadataDb(txCtx).DeleteAll() if err != nil { log.Error("error reest collection metadata db", zap.Error(err)) return err } - err = tc.metaDomain.SegmentDb(txCtx).DeleteAll() + err = tc.metaDomain.CollectionDb(txCtx).DeleteAll() if err != nil { - log.Error("error reset segment db", zap.Error(err)) + log.Error("error reset collection db", zap.Error(err)) return err } err = tc.metaDomain.SegmentMetadataDb(txCtx).DeleteAll() @@ -57,12 +53,19 @@ func (tc *Catalog) ResetState(ctx context.Context) error { log.Error("error reset segment metadata db", zap.Error(err)) return err } + err = tc.metaDomain.SegmentDb(txCtx).DeleteAll() + if err != nil { + log.Error("error reset segment db", zap.Error(err)) + return err + } + err = tc.metaDomain.DatabaseDb(txCtx).DeleteAll() if err != nil { log.Error("error reset database db", zap.Error(err)) return err } + // TODO: default database and tenant should be pre-defined object err = tc.metaDomain.DatabaseDb(txCtx).Insert(&dbmodel.Database{ ID: types.NilUniqueID().String(), Name: common.DefaultDatabase, @@ -79,7 +82,8 @@ func (tc *Catalog) ResetState(ctx context.Context) error { return err } err = tc.metaDomain.TenantDb(txCtx).Insert(&dbmodel.Tenant{ - ID: common.DefaultTenant, + ID: common.DefaultTenant, + LastCompactionTime: time.Now().Unix(), }) if err != nil { log.Error("error inserting default tenant", zap.Error(err)) @@ -153,9 +157,11 @@ func (tc *Catalog) CreateTenant(ctx context.Context, createTenant *model.CreateT var result *model.Tenant err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { + // TODO: createTenant has ts, don't need to pass in dbTenant := &dbmodel.Tenant{ - ID: createTenant.Name, - Ts: ts, + ID: createTenant.Name, + Ts: ts, + LastCompactionTime: time.Now().Unix(), } err := tc.metaDomain.TenantDb(txCtx).Insert(dbTenant) if err != nil { @@ -222,7 +228,7 @@ func (tc *Catalog) CreateCollection(ctx context.Context, createCollection *model } collectionName := createCollection.Name - existing, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(types.FromUniqueID(createCollection.ID), &collectionName, nil, tenantID, databaseName) + existing, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(nil, &collectionName, tenantID, databaseName) if err != nil { log.Error("error getting collection", zap.Error(err)) return err @@ -251,12 +257,12 @@ func (tc *Catalog) CreateCollection(ctx context.Context, createCollection *model } dbCollection := &dbmodel.Collection{ - ID: createCollection.ID.String(), - Name: &createCollection.Name, - Topic: &createCollection.Topic, - Dimension: createCollection.Dimension, - DatabaseID: databases[0].ID, - Ts: ts, + ID: createCollection.ID.String(), + Name: &createCollection.Name, + Dimension: createCollection.Dimension, + DatabaseID: databases[0].ID, + Ts: ts, + LogPosition: 0, } err = tc.metaDomain.CollectionDb(txCtx).Insert(dbCollection) @@ -274,13 +280,12 @@ func (tc *Catalog) CreateCollection(ctx context.Context, createCollection *model } } // get collection - collectionList, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(types.FromUniqueID(createCollection.ID), nil, nil, tenantID, databaseName) + collectionList, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(types.FromUniqueID(createCollection.ID), nil, tenantID, databaseName) if err != nil { log.Error("error getting collection", zap.Error(err)) return err } result = convertCollectionToModel(collectionList)[0] - result.Created = true notificationRecord := &dbmodel.Notification{ CollectionID: result.ID.String(), @@ -301,8 +306,8 @@ func (tc *Catalog) CreateCollection(ctx context.Context, createCollection *model return result, nil } -func (tc *Catalog) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenandID string, databaseName string) ([]*model.Collection, error) { - collectionAndMetadataList, err := tc.metaDomain.CollectionDb(ctx).GetCollections(types.FromUniqueID(collectionID), collectionName, collectionTopic, tenandID, databaseName) +func (tc *Catalog) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, tenantID string, databaseName string) ([]*model.Collection, error) { + collectionAndMetadataList, err := tc.metaDomain.CollectionDb(ctx).GetCollections(types.FromUniqueID(collectionID), collectionName, tenantID, databaseName) if err != nil { return nil, err } @@ -311,16 +316,27 @@ func (tc *Catalog) GetCollections(ctx context.Context, collectionID types.Unique } func (tc *Catalog) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error { + log.Info("deleting collection", zap.Any("deleteCollection", deleteCollection)) return tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { collectionID := deleteCollection.ID - err := tc.metaDomain.CollectionDb(txCtx).DeleteCollectionByID(collectionID.String()) + collectionAndMetadata, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(types.FromUniqueID(collectionID), nil, deleteCollection.TenantID, deleteCollection.DatabaseName) if err != nil { return err } - err = tc.metaDomain.CollectionMetadataDb(txCtx).DeleteByCollectionID(collectionID.String()) + if len(collectionAndMetadata) == 0 { + return common.ErrCollectionDeleteNonExistingCollection + } + + collectionDeletedCount, err := tc.metaDomain.CollectionDb(txCtx).DeleteCollectionByID(collectionID.String()) if err != nil { return err } + collectionMetadataDeletedCount, err := tc.metaDomain.CollectionMetadataDb(txCtx).DeleteByCollectionID(collectionID.String()) + if err != nil { + return err + } + log.Info("collection deleted", zap.Any("collection", collectionAndMetadata), zap.Int("collectionDeletedCount", collectionDeletedCount), zap.Int("collectionMetadataDeletedCount", collectionMetadataDeletedCount)) + notificationRecord := &dbmodel.Notification{ CollectionID: collectionID.String(), Type: dbmodel.NotificationTypeDeleteCollection, @@ -330,18 +346,19 @@ func (tc *Catalog) DeleteCollection(ctx context.Context, deleteCollection *model if err != nil { return err } + return nil }) } func (tc *Catalog) UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection, ts types.Timestamp) (*model.Collection, error) { + log.Info("updating collection", zap.String("collectionId", updateCollection.ID.String())) var result *model.Collection err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { dbCollection := &dbmodel.Collection{ ID: updateCollection.ID.String(), Name: updateCollection.Name, - Topic: updateCollection.Topic, Dimension: updateCollection.Dimension, Ts: ts, } @@ -360,14 +377,14 @@ func (tc *Catalog) UpdateCollection(ctx context.Context, updateCollection *model if metadata != nil { // Case 2 return common.ErrInvalidMetadataUpdate } else { // Case 1 - err = tc.metaDomain.CollectionMetadataDb(txCtx).DeleteByCollectionID(updateCollection.ID.String()) + _, err = tc.metaDomain.CollectionMetadataDb(txCtx).DeleteByCollectionID(updateCollection.ID.String()) if err != nil { return err } } } else { if metadata != nil { // Case 3 - err = tc.metaDomain.CollectionMetadataDb(txCtx).DeleteByCollectionID(updateCollection.ID.String()) + _, err = tc.metaDomain.CollectionMetadataDb(txCtx).DeleteByCollectionID(updateCollection.ID.String()) if err != nil { return err } @@ -382,17 +399,20 @@ func (tc *Catalog) UpdateCollection(ctx context.Context, updateCollection *model } databaseName := updateCollection.DatabaseName tenantID := updateCollection.TenantID - collectionList, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(types.FromUniqueID(updateCollection.ID), nil, nil, tenantID, databaseName) + collectionList, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(types.FromUniqueID(updateCollection.ID), nil, tenantID, databaseName) if err != nil { return err } + if collectionList == nil || len(collectionList) == 0 { + return common.ErrCollectionNotFound + } result = convertCollectionToModel(collectionList)[0] return nil }) if err != nil { return nil, err } - log.Info("collection updated", zap.Any("collection", result)) + log.Info("collection updated", zap.String("collectionID", result.ID.String())) return result, nil } @@ -409,9 +429,6 @@ func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.Creat Scope: createSegment.Scope, Ts: ts, } - if createSegment.Topic != nil { - dbSegment.Topic = createSegment.Topic - } err := tc.metaDomain.SegmentDb(txCtx).Insert(dbSegment) if err != nil { log.Error("error inserting segment", zap.Error(err)) @@ -430,7 +447,7 @@ func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.Creat } } // get segment - segmentList, err := tc.metaDomain.SegmentDb(txCtx).GetSegments(createSegment.ID, nil, nil, nil, types.NilUniqueID()) + segmentList, err := tc.metaDomain.SegmentDb(txCtx).GetSegments(createSegment.ID, nil, nil, types.NilUniqueID()) if err != nil { log.Error("error getting segment", zap.Error(err)) return err @@ -446,19 +463,19 @@ func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.Creat return result, nil } -func (tc *Catalog) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID, ts types.Timestamp) ([]*model.Segment, error) { - segmentAndMetadataList, err := tc.metaDomain.SegmentDb(ctx).GetSegments(segmentID, segmentType, scope, topic, collectionID) +func (tc *Catalog) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*model.Segment, error) { + segmentAndMetadataList, err := tc.metaDomain.SegmentDb(ctx).GetSegments(segmentID, segmentType, scope, collectionID) if err != nil { return nil, err } segments := make([]*model.Segment, 0, len(segmentAndMetadataList)) for _, segmentAndMetadata := range segmentAndMetadataList { segment := &model.Segment{ - ID: types.MustParse(segmentAndMetadata.Segment.ID), - Type: segmentAndMetadata.Segment.Type, - Scope: segmentAndMetadata.Segment.Scope, - Topic: segmentAndMetadata.Segment.Topic, - Ts: segmentAndMetadata.Segment.Ts, + ID: types.MustParse(segmentAndMetadata.Segment.ID), + Type: segmentAndMetadata.Segment.Type, + Scope: segmentAndMetadata.Segment.Scope, + Ts: segmentAndMetadata.Segment.Ts, + FilePaths: segmentAndMetadata.Segment.FilePaths, } if segmentAndMetadata.Segment.CollectionID != nil { @@ -474,7 +491,15 @@ func (tc *Catalog) GetSegments(ctx context.Context, segmentID types.UniqueID, se func (tc *Catalog) DeleteSegment(ctx context.Context, segmentID types.UniqueID) error { return tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { - err := tc.metaDomain.SegmentDb(txCtx).DeleteSegmentByID(segmentID.String()) + segment, err := tc.metaDomain.SegmentDb(txCtx).GetSegments(segmentID, nil, nil, types.NilUniqueID()) + if err != nil { + return err + } + if len(segment) == 0 { + return common.ErrSegmentDeleteNonExistingSegment + } + + err = tc.metaDomain.SegmentDb(txCtx).DeleteSegmentByID(segmentID.String()) if err != nil { log.Error("error deleting segment", zap.Error(err)) return err @@ -492,11 +517,25 @@ func (tc *Catalog) UpdateSegment(ctx context.Context, updateSegment *model.Updat var result *model.Segment err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { + // TODO: we should push in collection_id here, add a GET to fix test for now + if updateSegment.Collection == nil { + results, err := tc.metaDomain.SegmentDb(txCtx).GetSegments(updateSegment.ID, nil, nil, types.NilUniqueID()) + if err != nil { + return err + } + if results == nil || len(results) == 0 { + return common.ErrSegmentUpdateNonExistingSegment + } + if results != nil && len(results) > 1 { + // TODO: fix this error + return common.ErrInvalidCollectionUpdate + } + updateSegment.Collection = results[0].Segment.CollectionID + } + // update segment dbSegment := &dbmodel.UpdateSegment{ ID: updateSegment.ID.String(), - Topic: updateSegment.Topic, - ResetTopic: updateSegment.ResetTopic, Collection: updateSegment.Collection, ResetCollection: updateSegment.ResetCollection, } @@ -547,7 +586,7 @@ func (tc *Catalog) UpdateSegment(ctx context.Context, updateSegment *model.Updat } // get segment - segmentList, err := tc.metaDomain.SegmentDb(txCtx).GetSegments(updateSegment.ID, nil, nil, nil, types.NilUniqueID()) + segmentList, err := tc.metaDomain.SegmentDb(txCtx).GetSegments(updateSegment.ID, nil, nil, types.NilUniqueID()) if err != nil { log.Error("error getting segment", zap.Error(err)) return err @@ -562,3 +601,50 @@ func (tc *Catalog) UpdateSegment(ctx context.Context, updateSegment *model.Updat log.Debug("segment updated", zap.Any("segment", result)) return result, nil } + +func (tc *Catalog) SetTenantLastCompactionTime(ctx context.Context, tenantID string, lastCompactionTime int64) error { + return tc.metaDomain.TenantDb(ctx).UpdateTenantLastCompactionTime(tenantID, lastCompactionTime) +} + +func (tc *Catalog) GetTenantsLastCompactionTime(ctx context.Context, tenantIDs []string) ([]*dbmodel.Tenant, error) { + tenants, err := tc.metaDomain.TenantDb(ctx).GetTenantsLastCompactionTime(tenantIDs) + return tenants, err +} + +func (tc *Catalog) FlushCollectionCompaction(ctx context.Context, flushCollectionCompaction *model.FlushCollectionCompaction) (*model.FlushCollectionInfo, error) { + flushCollectionInfo := &model.FlushCollectionInfo{ + ID: flushCollectionCompaction.ID.String(), + } + + err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error { + // register files to Segment metadata + err := tc.metaDomain.SegmentDb(txCtx).RegisterFilePaths(flushCollectionCompaction.FlushSegmentCompactions) + if err != nil { + return err + } + + // update collection log position and version + collectionVersion, err := tc.metaDomain.CollectionDb(txCtx).UpdateLogPositionAndVersion(flushCollectionCompaction.ID.String(), flushCollectionCompaction.LogPosition, flushCollectionCompaction.CurrentCollectionVersion) + if err != nil { + return err + } + flushCollectionInfo.CollectionVersion = collectionVersion + + // update tenant last compaction time + // TODO: add a system configuration to disable + // since this might cause resource contention if one tenant has a lot of collection compactions at the same time + lastCompactionTime := time.Now().Unix() + err = tc.metaDomain.TenantDb(txCtx).UpdateTenantLastCompactionTime(flushCollectionCompaction.TenantID, lastCompactionTime) + if err != nil { + return err + } + flushCollectionInfo.TenantLastCompactionTime = lastCompactionTime + + // return nil will commit the transaction + return nil + }) + if err != nil { + return nil, err + } + return flushCollectionInfo, nil +} diff --git a/go/coordinator/internal/metastore/coordinator/table_catalog_test.go b/go/pkg/metastore/coordinator/table_catalog_test.go similarity index 83% rename from go/coordinator/internal/metastore/coordinator/table_catalog_test.go rename to go/pkg/metastore/coordinator/table_catalog_test.go index f40cddffd38..d913925550e 100644 --- a/go/coordinator/internal/metastore/coordinator/table_catalog_test.go +++ b/go/pkg/metastore/coordinator/table_catalog_test.go @@ -4,15 +4,20 @@ import ( "context" "testing" - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel/mocks" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel/mocks" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" ) +const ( + defaultTenant = "default_tenant" + defaultDatabase = "default_database" +) + func TestCatalog_CreateCollection(t *testing.T) { // create a mock transaction implementation mockTxImpl := &mocks.ITransaction{} @@ -42,8 +47,7 @@ func TestCatalog_CreateCollection(t *testing.T) { mockMetaDomain.CollectionDb(context.Background()).(*mocks.ICollectionDb).On("Insert", &dbmodel.Collection{ ID: "00000000-0000-0000-0000-000000000001", Name: &name, - // Topic: "test_topic", - Ts: ts, + Ts: ts, }).Return(nil) // mock the insert collection metadata method @@ -82,9 +86,6 @@ func TestCatalog_GetCollections(t *testing.T) { // create a mock collection name collectionName := "test_collection" - // create a mock collection topic - collectionTopic := "test_topic" - // create a mock collection and metadata list name := "test_collection" testKey := "test_key" @@ -92,10 +93,9 @@ func TestCatalog_GetCollections(t *testing.T) { collectionAndMetadataList := []*dbmodel.CollectionAndMetadata{ { Collection: &dbmodel.Collection{ - ID: "00000000-0000-0000-0000-000000000001", - Name: &name, - Topic: &collectionTopic, - Ts: types.Timestamp(1234567890), + ID: "00000000-0000-0000-0000-000000000001", + Name: &name, + Ts: types.Timestamp(1234567890), }, CollectionMetadata: []*dbmodel.CollectionMetadata{ { @@ -110,10 +110,10 @@ func TestCatalog_GetCollections(t *testing.T) { // mock the get collections method mockMetaDomain.On("CollectionDb", context.Background()).Return(&mocks.ICollectionDb{}) - mockMetaDomain.CollectionDb(context.Background()).(*mocks.ICollectionDb).On("GetCollections", types.FromUniqueID(collectionID), &collectionName, &collectionTopic, common.DefaultTenant, common.DefaultDatabase).Return(collectionAndMetadataList, nil) + mockMetaDomain.CollectionDb(context.Background()).(*mocks.ICollectionDb).On("GetCollections", types.FromUniqueID(collectionID), &collectionName, common.DefaultTenant, common.DefaultDatabase).Return(collectionAndMetadataList, nil) // call the GetCollections method - collections, err := catalog.GetCollections(context.Background(), collectionID, &collectionName, &collectionTopic, defaultTenant, defaultDatabase) + collections, err := catalog.GetCollections(context.Background(), collectionID, &collectionName, defaultTenant, defaultDatabase) // assert that the method returned no error assert.NoError(t, err) @@ -125,7 +125,6 @@ func TestCatalog_GetCollections(t *testing.T) { { ID: types.MustParse("00000000-0000-0000-0000-000000000001"), Name: "test_collection", - Topic: collectionTopic, Ts: types.Timestamp(1234567890), Metadata: metadata, }, diff --git a/go/coordinator/internal/metastore/db/dao/collection.go b/go/pkg/metastore/db/dao/collection.go similarity index 51% rename from go/coordinator/internal/metastore/db/dao/collection.go rename to go/pkg/metastore/db/dao/collection.go index b21da7a9f76..3a41b833022 100644 --- a/go/coordinator/internal/metastore/db/dao/collection.go +++ b/go/pkg/metastore/db/dao/collection.go @@ -2,11 +2,17 @@ package dao import ( "database/sql" + "errors" + "strings" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/jackc/pgx/v5/pgconn" + "gorm.io/gorm/clause" "go.uber.org/zap" "gorm.io/gorm" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" "github.com/pingcap/log" ) @@ -20,28 +26,37 @@ func (s *collectionDb) DeleteAll() error { return s.db.Where("1 = 1").Delete(&dbmodel.Collection{}).Error } -func (s *collectionDb) GetCollections(id *string, name *string, topic *string, tenantID string, databaseName string) ([]*dbmodel.CollectionAndMetadata, error) { +func (s *collectionDb) GetCollections(id *string, name *string, tenantID string, databaseName string) ([]*dbmodel.CollectionAndMetadata, error) { + var getCollectionInput strings.Builder + getCollectionInput.WriteString("GetCollections input: ") + var collections []*dbmodel.CollectionAndMetadata query := s.db.Table("collections"). - Select("collections.id, collections.name, collections.topic, collections.dimension, collections.database_id, databases.name, databases.tenant_id, collection_metadata.key, collection_metadata.str_value, collection_metadata.int_value, collection_metadata.float_value"). + Select("collections.id, collections.log_position, collections.version, collections.name, collections.dimension, collections.database_id, databases.name, databases.tenant_id, collection_metadata.key, collection_metadata.str_value, collection_metadata.int_value, collection_metadata.float_value"). Joins("LEFT JOIN collection_metadata ON collections.id = collection_metadata.collection_id"). Joins("INNER JOIN databases ON collections.database_id = databases.id"). Order("collections.id") - query = query.Where("databases.name = ?", databaseName) + if databaseName != "" { + query = query.Where("databases.name = ?", databaseName) + getCollectionInput.WriteString("databases.name: " + databaseName + ", ") + } - query = query.Where("databases.tenant_id = ?", tenantID) + if tenantID != "" { + query = query.Where("databases.tenant_id = ?", tenantID) + getCollectionInput.WriteString("databases.tenant_id: " + tenantID + ", ") + } if id != nil { query = query.Where("collections.id = ?", *id) - } - if topic != nil { - query = query.Where("collections.topic = ?", *topic) + getCollectionInput.WriteString("collections.id: " + *id + ", ") } if name != nil { query = query.Where("collections.name = ?", *name) + getCollectionInput.WriteString("collections.name: " + *name + ", ") } + log.Info(getCollectionInput.String()) rows, err := query.Rows() if err != nil { @@ -56,8 +71,9 @@ func (s *collectionDb) GetCollections(id *string, name *string, topic *string, t for rows.Next() { var ( collectionID string + logPosition int64 + version int32 collectionName string - collectionTopic string collectionDimension sql.NullInt32 collectionDatabaseID string databaseName string @@ -68,7 +84,7 @@ func (s *collectionDb) GetCollections(id *string, name *string, topic *string, t floatValue sql.NullFloat64 ) - err := rows.Scan(&collectionID, &collectionName, &collectionTopic, &collectionDimension, &collectionDatabaseID, &databaseName, &databaseTenantID, &key, &strValue, &intValue, &floatValue) + err := rows.Scan(&collectionID, &logPosition, &version, &collectionName, &collectionDimension, &collectionDatabaseID, &databaseName, &databaseTenantID, &key, &strValue, &intValue, &floatValue) if err != nil { log.Error("scan collection failed", zap.Error(err)) return nil, err @@ -79,10 +95,11 @@ func (s *collectionDb) GetCollections(id *string, name *string, topic *string, t currentCollection = &dbmodel.CollectionAndMetadata{ Collection: &dbmodel.Collection{ - ID: collectionID, - Name: &collectionName, - Topic: &collectionTopic, - DatabaseID: collectionDatabaseID, + ID: collectionID, + Name: &collectionName, + DatabaseID: collectionDatabaseID, + LogPosition: logPosition, + Version: version, }, CollectionMetadata: metadata, TenantID: databaseTenantID, @@ -132,12 +149,31 @@ func (s *collectionDb) GetCollections(id *string, name *string, topic *string, t return collections, nil } -func (s *collectionDb) DeleteCollectionByID(collectionID string) error { - return s.db.Where("id = ?", collectionID).Delete(&dbmodel.Collection{}).Error +func (s *collectionDb) DeleteCollectionByID(collectionID string) (int, error) { + var collections []dbmodel.Collection + err := s.db.Clauses(clause.Returning{}).Where("id = ?", collectionID).Delete(&collections).Error + return len(collections), err } func (s *collectionDb) Insert(in *dbmodel.Collection) error { - return s.db.Create(&in).Error + err := s.db.Create(&in).Error + if err != nil { + log.Error("create collection failed", zap.Error(err)) + var pgErr *pgconn.PgError + ok := errors.As(err, &pgErr) + if ok { + log.Error("Postgres Error") + switch pgErr.Code { + case "23505": + log.Error("collection already exists") + return common.ErrCollectionUniqueConstraintViolation + default: + return err + } + } + return err + } + return nil } func generateCollectionUpdatesWithoutID(in *dbmodel.Collection) map[string]interface{} { @@ -145,9 +181,6 @@ func generateCollectionUpdatesWithoutID(in *dbmodel.Collection) map[string]inter if in.Name != nil { ret["name"] = *in.Name } - if in.Topic != nil { - ret["topic"] = *in.Topic - } if in.Dimension != nil { ret["dimension"] = *in.Dimension } @@ -155,6 +188,33 @@ func generateCollectionUpdatesWithoutID(in *dbmodel.Collection) map[string]inter } func (s *collectionDb) Update(in *dbmodel.Collection) error { + log.Info("update collection", zap.Any("collection", in)) updates := generateCollectionUpdatesWithoutID(in) return s.db.Model(&dbmodel.Collection{}).Where("id = ?", in.ID).Updates(updates).Error } + +func (s *collectionDb) UpdateLogPositionAndVersion(collectionID string, logPosition int64, currentCollectionVersion int32) (int32, error) { + log.Info("update log position and version", zap.String("collectionID", collectionID), zap.Int64("logPosition", logPosition), zap.Int32("currentCollectionVersion", currentCollectionVersion)) + var collection dbmodel.Collection + err := s.db.Where("id = ?", collectionID).First(&collection).Error + if err != nil { + return 0, err + } + if collection.LogPosition > logPosition { + return 0, common.ErrCollectionLogPositionStale + } + if collection.Version > currentCollectionVersion { + return 0, common.ErrCollectionVersionStale + } + if collection.Version < currentCollectionVersion { + // this should not happen, potentially a bug + return 0, common.ErrCollectionVersionInvalid + } + + version := currentCollectionVersion + 1 + err = s.db.Model(&dbmodel.Collection{}).Where("id = ?", collectionID).Updates(map[string]interface{}{"log_position": logPosition, "version": version}).Error + if err != nil { + return 0, err + } + return version, nil +} diff --git a/go/coordinator/internal/metastore/db/dao/collection_metadata.go b/go/pkg/metastore/db/dao/collection_metadata.go similarity index 70% rename from go/coordinator/internal/metastore/db/dao/collection_metadata.go rename to go/pkg/metastore/db/dao/collection_metadata.go index 0f9ba00057e..c5790e94d38 100644 --- a/go/coordinator/internal/metastore/db/dao/collection_metadata.go +++ b/go/pkg/metastore/db/dao/collection_metadata.go @@ -1,7 +1,7 @@ package dao import ( - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" "gorm.io/gorm" "gorm.io/gorm/clause" ) @@ -14,8 +14,10 @@ func (s *collectionMetadataDb) DeleteAll() error { return s.db.Where("1 = 1").Delete(&dbmodel.CollectionMetadata{}).Error } -func (s *collectionMetadataDb) DeleteByCollectionID(collectionID string) error { - return s.db.Where("collection_id = ?", collectionID).Delete(&dbmodel.CollectionMetadata{}).Error +func (s *collectionMetadataDb) DeleteByCollectionID(collectionID string) (int, error) { + var metadata []dbmodel.CollectionMetadata + err := s.db.Clauses(clause.Returning{}).Where("collection_id = ?", collectionID).Delete(&metadata).Error + return len(metadata), err } func (s *collectionMetadataDb) Insert(in []*dbmodel.CollectionMetadata) error { diff --git a/go/pkg/metastore/db/dao/collection_test.go b/go/pkg/metastore/db/dao/collection_test.go new file mode 100644 index 00000000000..7be7828e225 --- /dev/null +++ b/go/pkg/metastore/db/dao/collection_test.go @@ -0,0 +1,131 @@ +package dao + +import ( + "testing" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "gorm.io/gorm" +) + +type CollectionDbTestSuite struct { + suite.Suite + db *gorm.DB + collectionDb *collectionDb + tenantName string + databaseName string + databaseId string +} + +func (suite *CollectionDbTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() + suite.collectionDb = &collectionDb{ + db: suite.db, + } + suite.tenantName = "test_collection_tenant" + suite.databaseName = "test_collection_database" + DbId, err := CreateTestTenantAndDatabase(suite.db, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.databaseId = DbId +} + +func (suite *CollectionDbTestSuite) TearDownSuite() { + log.Info("teardown suite") + err := CleanUpTestDatabase(suite.db, suite.tenantName, suite.databaseName) + suite.NoError(err) + err = CleanUpTestTenant(suite.db, suite.tenantName) + suite.NoError(err) +} + +func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() { + collectionName := "test_collection_get_collections" + collectionID, err := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId) + suite.NoError(err) + + testKey := "test" + testValue := "test" + metadata := &dbmodel.CollectionMetadata{ + CollectionID: collectionID, + Key: &testKey, + StrValue: &testValue, + } + err = suite.db.Create(metadata).Error + suite.NoError(err) + + query := suite.db.Table("collections").Select("collections.id").Where("collections.id = ?", collectionID) + rows, err := query.Rows() + suite.NoError(err) + for rows.Next() { + var scanedCollectionID string + err = rows.Scan(&scanedCollectionID) + suite.NoError(err) + suite.Equal(collectionID, scanedCollectionID) + } + collections, err := suite.collectionDb.GetCollections(nil, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Len(collections, 1) + suite.Equal(collectionID, collections[0].Collection.ID) + suite.Equal(collectionName, *collections[0].Collection.Name) + suite.Len(collections[0].CollectionMetadata, 1) + suite.Equal(metadata.Key, collections[0].CollectionMetadata[0].Key) + suite.Equal(metadata.StrValue, collections[0].CollectionMetadata[0].StrValue) + + // Test when filtering by ID + collections, err = suite.collectionDb.GetCollections(nil, nil, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Len(collections, 1) + suite.Equal(collectionID, collections[0].Collection.ID) + + // Test when filtering by name + collections, err = suite.collectionDb.GetCollections(nil, &collectionName, suite.tenantName, suite.databaseName) + suite.NoError(err) + suite.Len(collections, 1) + suite.Equal(collectionID, collections[0].Collection.ID) + + // clean up + err = CleanUpTestCollection(suite.db, collectionID) + suite.NoError(err) +} + +func (suite *CollectionDbTestSuite) TestCollectionDb_UpdateLogPositionAndVersion() { + collectionName := "test_collection_get_collections" + collectionID, err := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId) + // verify default values + collections, err := suite.collectionDb.GetCollections(&collectionID, nil, "", "") + suite.NoError(err) + suite.Len(collections, 1) + suite.Equal(int64(0), collections[0].Collection.LogPosition) + suite.Equal(int32(0), collections[0].Collection.Version) + + // update log position and version + version, err := suite.collectionDb.UpdateLogPositionAndVersion(collectionID, int64(10), 0) + suite.NoError(err) + suite.Equal(int32(1), version) + collections, err = suite.collectionDb.GetCollections(&collectionID, nil, "", "") + suite.Len(collections, 1) + suite.Equal(int64(10), collections[0].Collection.LogPosition) + suite.Equal(int32(1), collections[0].Collection.Version) + + // invalid log position + _, err = suite.collectionDb.UpdateLogPositionAndVersion(collectionID, int64(5), 0) + suite.Error(err, "collection log position Stale") + + // invalid version + _, err = suite.collectionDb.UpdateLogPositionAndVersion(collectionID, int64(20), 0) + suite.Error(err, "collection version invalid") + _, err = suite.collectionDb.UpdateLogPositionAndVersion(collectionID, int64(20), 3) + suite.Error(err, "collection version invalid") + + //clean up + err = CleanUpTestCollection(suite.db, collectionID) + suite.NoError(err) +} + +func TestCollectionDbTestSuiteSuite(t *testing.T) { + testSuite := new(CollectionDbTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/coordinator/internal/metastore/db/dao/common.go b/go/pkg/metastore/db/dao/common.go similarity index 80% rename from go/coordinator/internal/metastore/db/dao/common.go rename to go/pkg/metastore/db/dao/common.go index c67cea6c759..ff30e2b09d6 100644 --- a/go/coordinator/internal/metastore/db/dao/common.go +++ b/go/pkg/metastore/db/dao/common.go @@ -3,8 +3,8 @@ package dao import ( "context" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbcore" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" ) type metaDomain struct{} @@ -40,3 +40,7 @@ func (*metaDomain) SegmentMetadataDb(ctx context.Context) dbmodel.ISegmentMetada func (*metaDomain) NotificationDb(ctx context.Context) dbmodel.INotificationDb { return ¬ificationDb{dbcore.GetDB(ctx)} } + +func (*metaDomain) RecordLogDb(ctx context.Context) dbmodel.IRecordLogDb { + return &recordLogDb{dbcore.GetDB(ctx)} +} diff --git a/go/coordinator/internal/metastore/db/dao/database.go b/go/pkg/metastore/db/dao/database.go similarity index 57% rename from go/coordinator/internal/metastore/db/dao/database.go rename to go/pkg/metastore/db/dao/database.go index 0d02dca484b..fb7ffb07a12 100644 --- a/go/coordinator/internal/metastore/db/dao/database.go +++ b/go/pkg/metastore/db/dao/database.go @@ -1,10 +1,11 @@ package dao import ( - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" "github.com/pingcap/log" "go.uber.org/zap" "gorm.io/gorm" + "gorm.io/gorm/clause" ) type databaseDb struct { @@ -17,6 +18,12 @@ func (s *databaseDb) DeleteAll() error { return s.db.Where("1 = 1").Delete(&dbmodel.Database{}).Error } +func (s *databaseDb) DeleteByTenantIdAndName(tenantId string, databaseName string) (int, error) { + var databases []dbmodel.Database + err := s.db.Clauses(clause.Returning{}).Where("tenant_id = ?", tenantId).Where("name = ?", databaseName).Delete(&databases).Error + return len(databases), err +} + func (s *databaseDb) GetAllDatabases() ([]*dbmodel.Database, error) { var databases []*dbmodel.Database query := s.db.Table("databases") @@ -44,3 +51,16 @@ func (s *databaseDb) GetDatabases(tenantID string, databaseName string) ([]*dbmo func (s *databaseDb) Insert(database *dbmodel.Database) error { return s.db.Create(database).Error } + +func (s *databaseDb) GetDatabasesByTenantID(tenantID string) ([]*dbmodel.Database, error) { + var databases []*dbmodel.Database + query := s.db.Table("databases"). + Select("databases.id, databases.name, databases.tenant_id"). + Where("databases.tenant_id = ?", tenantID) + + if err := query.Find(&databases).Error; err != nil { + log.Error("GetDatabasesByTenantID", zap.Error(err)) + return nil, err + } + return databases, nil +} diff --git a/go/coordinator/internal/metastore/db/dao/notification.go b/go/pkg/metastore/db/dao/notification.go similarity index 94% rename from go/coordinator/internal/metastore/db/dao/notification.go rename to go/pkg/metastore/db/dao/notification.go index cfd6bfdfbbf..2b5e2231554 100644 --- a/go/coordinator/internal/metastore/db/dao/notification.go +++ b/go/pkg/metastore/db/dao/notification.go @@ -1,7 +1,7 @@ package dao import ( - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" "gorm.io/gorm" ) diff --git a/go/pkg/metastore/db/dao/record_log.go b/go/pkg/metastore/db/dao/record_log.go new file mode 100644 index 00000000000..c967f5ebf7d --- /dev/null +++ b/go/pkg/metastore/db/dao/record_log.go @@ -0,0 +1,131 @@ +package dao + +import ( + "database/sql" + "errors" + "time" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/pingcap/log" + "go.uber.org/zap" + "gorm.io/gorm" +) + +type recordLogDb struct { + db *gorm.DB +} + +var _ dbmodel.IRecordLogDb = &recordLogDb{} + +func (s *recordLogDb) PushLogs(collectionID types.UniqueID, recordsContent [][]byte) (int, error) { + err := s.db.Transaction(func(tx *gorm.DB) error { + var timestamp = time.Now().UnixNano() + var collectionIDStr = types.FromUniqueID(collectionID) + log.Info("PushLogs", + zap.String("collectionID", *collectionIDStr), + zap.Int64("timestamp", timestamp), + zap.Int("count", len(recordsContent))) + + var lastLog *dbmodel.RecordLog + err := tx.Select("log_offset").Where("collection_id = ?", collectionIDStr).Order("log_offset DESC").Limit(1).Find(&lastLog).Error + if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { + log.Error("Get last log offset error", zap.Error(err)) + tx.Rollback() + return err + } + // The select will populate the lastLog with the last log in the collection, if + // one does not exist, it will have a default value of 0, so we can safely use it + var lastLogOffset = lastLog.LogOffset + log.Info("PushLogs", zap.Int64("lastLogOffset", lastLogOffset)) + + var recordLogs []*dbmodel.RecordLog + for index := range recordsContent { + recordLogs = append(recordLogs, &dbmodel.RecordLog{ + CollectionID: collectionIDStr, + LogOffset: lastLogOffset + int64(index) + 1, + Timestamp: timestamp, + Record: &recordsContent[index], + }) + } + err = tx.CreateInBatches(recordLogs, len(recordLogs)).Error + if err != nil { + log.Error("Batch insert error", zap.Error(err)) + tx.Rollback() + return err + } + return nil + }) + if err != nil { + log.Error("PushLogs error", zap.Error(err)) + return 0, err + } + return len(recordsContent), nil +} + +func (s *recordLogDb) PullLogs(collectionID types.UniqueID, offset int64, batchSize int, endTimestamp int64) ([]*dbmodel.RecordLog, error) { + var collectionIDStr = types.FromUniqueID(collectionID) + log.Info("PullLogs", + zap.String("collectionID", *collectionIDStr), + zap.Int64("log_offset", offset), + zap.Int("batch_size", batchSize), + zap.Int64("endTimestamp", endTimestamp)) + + var recordLogs []*dbmodel.RecordLog + if endTimestamp > 0 { + result := s.db.Where("collection_id = ? AND log_offset >= ? AND timestamp <= ?", collectionIDStr, offset, endTimestamp).Order("log_offset").Limit(batchSize).Find(&recordLogs) + if result.Error != nil && !errors.Is(result.Error, gorm.ErrRecordNotFound) { + log.Error("PullLogs error", zap.Error(result.Error)) + return nil, result.Error + } + } else { + result := s.db.Where("collection_id = ? AND log_offset >= ?", collectionIDStr, offset).Order("log_offset").Limit(batchSize).Find(&recordLogs) + if result.Error != nil && !errors.Is(result.Error, gorm.ErrRecordNotFound) { + log.Error("PullLogs error", zap.Error(result.Error)) + return nil, result.Error + } + } + log.Info("PullLogs", + zap.String("collectionID", *collectionIDStr), + zap.Int64("log_offset", offset), + zap.Int("batch_size", batchSize), + zap.Int("count", len(recordLogs))) + return recordLogs, nil +} + +func (s *recordLogDb) GetAllCollectionsToCompact() ([]*dbmodel.RecordLog, error) { + log.Info("GetAllCollectionsToCompact") + var recordLogs []*dbmodel.RecordLog + var rawSql = ` + with summary as ( + select r.collection_id, r.log_offset, r.timestamp, row_number() over(partition by r.collection_id order by r.log_offset) as rank + from record_logs r, collections c + where r.collection_id = c.id + and r.log_offset>c.log_position + ) + select * from summary + where rank=1 + order by timestamp;` + rows, err := s.db.Raw(rawSql).Rows() + defer func(rows *sql.Rows) { + err := rows.Close() + if err != nil { + log.Error("GetAllCollectionsToCompact Close error", zap.Error(err)) + } + }(rows) + if err != nil { + log.Error("GetAllCollectionsToCompact error", zap.Error(err)) + return nil, err + } + for rows.Next() { + var batchRecordLogs []*dbmodel.RecordLog + err := s.db.ScanRows(rows, &recordLogs) + if err != nil { + log.Error("GetAllCollectionsToCompact ScanRows error", zap.Error(err)) + return nil, err + } + recordLogs = append(recordLogs, batchRecordLogs...) + } + log.Info("GetAllCollectionsToCompact find collections count", zap.Int("count", len(recordLogs))) + return recordLogs, nil +} diff --git a/go/pkg/metastore/db/dao/record_log_test.go b/go/pkg/metastore/db/dao/record_log_test.go new file mode 100644 index 00000000000..49f5019a0ae --- /dev/null +++ b/go/pkg/metastore/db/dao/record_log_test.go @@ -0,0 +1,198 @@ +package dao + +import ( + "testing" + "time" + + "github.com/chroma-core/chroma/go/pkg/logservice/testutils" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "gorm.io/gorm" +) + +type RecordLogDbTestSuite struct { + suite.Suite + db *gorm.DB + Db *recordLogDb + collectionId1 types.UniqueID + collectionId2 types.UniqueID + records [][]byte +} + +func (suite *RecordLogDbTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() + suite.Db = &recordLogDb{ + db: suite.db, + } + suite.records = make([][]byte, 0, 5) + suite.records = append(suite.records, []byte("test1"), []byte("test2"), + []byte("test3"), []byte("test4"), []byte("test5")) + recordLogTableExist := suite.db.Migrator().HasTable(&dbmodel.RecordLog{}) + if !recordLogTableExist { + err := suite.db.Migrator().CreateTable(&dbmodel.RecordLog{}) + suite.NoError(err) + } +} + +func (suite *RecordLogDbTestSuite) SetupTest() { + log.Info("setup test") + suite.collectionId1 = types.NewUniqueID() + suite.collectionId2 = types.NewUniqueID() + err := testutils.CreateCollections(suite.db, suite.collectionId1, suite.collectionId2) + suite.NoError(err) +} + +func (suite *RecordLogDbTestSuite) TearDownTest() { + log.Info("teardown test") + err := testutils.CleanupCollections(suite.db, suite.collectionId1, suite.collectionId2) + suite.NoError(err) +} + +func (suite *RecordLogDbTestSuite) TestRecordLogDb_PushLogs() { + // run push logs in transaction + // id: 0, + // records: test1, test2, test3 + count, err := suite.Db.PushLogs(suite.collectionId1, suite.records[:3]) + suite.NoError(err) + suite.Equal(3, count) + + // verify logs are pushed + var recordLogs []*dbmodel.RecordLog + suite.db.Where("collection_id = ?", types.FromUniqueID(suite.collectionId1)).Find(&recordLogs) + suite.Len(recordLogs, 3) + for index := range recordLogs { + suite.Equal(int64(index+1), recordLogs[index].LogOffset) + suite.Equal(suite.records[index], *recordLogs[index].Record) + } + + // run push logs in transaction + // id: 1, + // records: test4, test5 + count, err = suite.Db.PushLogs(suite.collectionId1, suite.records[3:]) + suite.NoError(err) + suite.Equal(2, count) + + // verify logs are pushed + suite.db.Where("collection_id = ?", types.FromUniqueID(suite.collectionId1)).Find(&recordLogs) + suite.Len(recordLogs, 5) + for index := range recordLogs { + suite.Equal(int64(index+1), recordLogs[index].LogOffset) + suite.Equal(suite.records[index], *recordLogs[index].Record) + } + + // run push logs in transaction + // id: 0, + // records: test1, test2, test3, test4, test5 + count, err = suite.Db.PushLogs(suite.collectionId2, suite.records) + suite.NoError(err) + suite.Equal(5, count) + + // verify logs are pushed + suite.db.Where("collection_id = ?", types.FromUniqueID(suite.collectionId2)).Find(&recordLogs) + suite.Len(recordLogs, 5) + for index := range recordLogs { + suite.Equal(int64(index+1), recordLogs[index].LogOffset) + suite.Equal(suite.records[index], *recordLogs[index].Record) + } +} + +func (suite *RecordLogDbTestSuite) TestRecordLogDb_PullLogsFromID() { + // pull empty logs + var recordLogs []*dbmodel.RecordLog + invalidEndTimestamp := int64(-1) + recordLogs, err := suite.Db.PullLogs(suite.collectionId1, 0, 3, invalidEndTimestamp) + suite.NoError(err) + suite.Len(recordLogs, 0) + + // push some logs + count, err := suite.Db.PushLogs(suite.collectionId1, suite.records[:3]) + suite.NoError(err) + suite.Equal(3, count) + count, err = suite.Db.PushLogs(suite.collectionId1, suite.records[3:]) + suite.NoError(err) + suite.Equal(2, count) + + // pull logs from id 0 batch_size 3 + recordLogs, err = suite.Db.PullLogs(suite.collectionId1, 0, 3, invalidEndTimestamp) + suite.NoError(err) + suite.Len(recordLogs, 3) + for index := range recordLogs { + suite.Equal(int64(index+1), recordLogs[index].LogOffset) + suite.Equal(suite.records[index], *recordLogs[index].Record) + } + + // pull logs from id 0 batch_size 6 + recordLogs, err = suite.Db.PullLogs(suite.collectionId1, 0, 6, invalidEndTimestamp) + suite.NoError(err) + suite.Len(recordLogs, 5) + + for index := range recordLogs { + suite.Equal(int64(index+1), recordLogs[index].LogOffset) + suite.Equal(suite.records[index], *recordLogs[index].Record) + } + + // pull logs from id 3 batch_size 4 + recordLogs, err = suite.Db.PullLogs(suite.collectionId1, 3, 4, invalidEndTimestamp) + suite.NoError(err) + suite.Len(recordLogs, 3) + for index := range recordLogs { + suite.Equal(int64(index+3), recordLogs[index].LogOffset) + suite.Equal(suite.records[index+2], *recordLogs[index].Record) + } + + // pull logs from id 3 batch_size 4 endTimestamp Now + recordLogs, err = suite.Db.PullLogs(suite.collectionId1, 3, 4, time.Now().UnixNano()) + suite.NoError(err) + suite.Len(recordLogs, 3) + for index := range recordLogs { + suite.Equal(int64(index+3), recordLogs[index].LogOffset) + suite.Equal(suite.records[index+2], *recordLogs[index].Record) + } +} + +func (suite *RecordLogDbTestSuite) TestRecordLogDb_GetAllCollectionsToCompact() { + // push some logs + count, err := suite.Db.PushLogs(suite.collectionId1, suite.records) + suite.NoError(err) + suite.Equal(5, count) + + // get all collection ids to compact + collectionInfos, err := suite.Db.GetAllCollectionsToCompact() + suite.NoError(err) + suite.Len(collectionInfos, 1) + suite.Equal(suite.collectionId1.String(), *collectionInfos[0].CollectionID) + suite.Equal(int64(1), collectionInfos[0].LogOffset) + + // move log position + testutils.MoveLogPosition(suite.db, suite.collectionId1, 2) + + // get all collection ids to compact + collectionInfos, err = suite.Db.GetAllCollectionsToCompact() + suite.NoError(err) + suite.Len(collectionInfos, 1) + suite.Equal(suite.collectionId1.String(), *collectionInfos[0].CollectionID) + suite.Equal(int64(3), collectionInfos[0].LogOffset) + + // push some logs + count, err = suite.Db.PushLogs(suite.collectionId2, suite.records) + suite.NoError(err) + suite.Equal(5, count) + + // get all collection ids to compact + collectionInfos, err = suite.Db.GetAllCollectionsToCompact() + suite.NoError(err) + suite.Len(collectionInfos, 2) + suite.Equal(suite.collectionId1.String(), *collectionInfos[0].CollectionID) + suite.Equal(int64(3), collectionInfos[0].LogOffset) + suite.Equal(suite.collectionId2.String(), *collectionInfos[1].CollectionID) + suite.Equal(int64(1), collectionInfos[1].LogOffset) +} + +func TestRecordLogDbTestSuite(t *testing.T) { + testSuite := new(RecordLogDbTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/coordinator/internal/metastore/db/dao/segment.go b/go/pkg/metastore/db/dao/segment.go similarity index 56% rename from go/coordinator/internal/metastore/db/dao/segment.go rename to go/pkg/metastore/db/dao/segment.go index c4c3842e278..670cddb8267 100644 --- a/go/coordinator/internal/metastore/db/dao/segment.go +++ b/go/pkg/metastore/db/dao/segment.go @@ -2,9 +2,15 @@ package dao import ( "database/sql" + "encoding/json" + "errors" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/jackc/pgx/v5/pgconn" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" "github.com/pingcap/log" "go.uber.org/zap" "gorm.io/gorm" @@ -26,18 +32,31 @@ func (s *segmentDb) Insert(in *dbmodel.Segment) error { err := s.db.Create(&in).Error if err != nil { - log.Error("insert segment failed", zap.String("segmentID", in.ID), zap.Int64("ts", in.Ts), zap.Error(err)) + log.Error("create segment failed", zap.Error(err)) + var pgErr *pgconn.PgError + ok := errors.As(err, &pgErr) + if ok { + log.Error("Postgres Error") + switch pgErr.Code { + case "23505": + log.Error("segment already exists") + return common.ErrSegmentUniqueConstraintViolation + default: + return err + } + } return err } + return nil return nil } -func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error) { +func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error) { var segments []*dbmodel.SegmentAndMetadata query := s.db.Table("segments"). - Select("segments.id, segments.collection_id, segments.type, segments.scope, segments.topic, segment_metadata.key, segment_metadata.str_value, segment_metadata.int_value, segment_metadata.float_value"). + Select("segments.id, segments.collection_id, segments.type, segments.scope, segments.file_paths, segment_metadata.key, segment_metadata.str_value, segment_metadata.int_value, segment_metadata.float_value"). Joins("LEFT JOIN segment_metadata ON segments.id = segment_metadata.segment_id"). Order("segments.id") @@ -50,16 +69,13 @@ func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *s if scope != nil { query = query.Where("scope = ?", scope) } - if topic != nil { - query = query.Where("topic = ?", topic) - } if collectionID != types.NilUniqueID() { query = query.Where("collection_id = ?", collectionID.String()) } rows, err := query.Rows() if err != nil { - log.Error("get segments failed", zap.String("segmentID", id.String()), zap.String("segmentType", *segmentType), zap.String("scope", *scope), zap.String("collectionTopic", *topic), zap.Error(err)) + log.Error("get segments failed", zap.String("segmentID", id.String()), zap.String("segmentType", *segmentType), zap.String("scope", *scope), zap.Error(err)) return nil, err } defer rows.Close() @@ -70,18 +86,18 @@ func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *s for rows.Next() { var ( - segmentID string - collectionID sql.NullString - segmentType string - scope string - topic sql.NullString - key sql.NullString - strValue sql.NullString - intValue sql.NullInt64 - floatValue sql.NullFloat64 + segmentID string + collectionID sql.NullString + segmentType string + scope string + filePathsJson string + key sql.NullString + strValue sql.NullString + intValue sql.NullInt64 + floatValue sql.NullFloat64 ) - err := rows.Scan(&segmentID, &collectionID, &segmentType, &scope, &topic, &key, &strValue, &intValue, &floatValue) + err := rows.Scan(&segmentID, &collectionID, &segmentType, &scope, &filePathsJson, &key, &strValue, &intValue, &floatValue) if err != nil { log.Error("scan segment failed", zap.Error(err)) } @@ -89,11 +105,17 @@ func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *s currentSegmentID = segmentID metadata = nil + var filePaths map[string][]string + err := json.Unmarshal([]byte(filePathsJson), &filePaths) + if err != nil { + return nil, err + } currentSegment = &dbmodel.SegmentAndMetadata{ Segment: &dbmodel.Segment{ - ID: segmentID, - Type: segmentType, - Scope: scope, + ID: segmentID, + Type: segmentType, + Scope: scope, + FilePaths: filePaths, }, SegmentMetadata: metadata, } @@ -103,12 +125,6 @@ func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *s currentSegment.Segment.CollectionID = nil } - if topic.Valid { - currentSegment.Segment.Topic = &topic.String - } else { - currentSegment.Segment.Topic = nil - } - if currentSegmentID != "" { segments = append(segments, currentSegment) } @@ -149,36 +165,45 @@ func (s *segmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *s } func generateSegmentUpdatesWithoutID(in *dbmodel.UpdateSegment) map[string]interface{} { - // Case 1: if ResetTopic is true and topic is nil, then set the topic to nil - // Case 2: if ResetTopic is true and topic is not nil -> THIS SHOULD NEVER HAPPEN - // Case 3: if ResetTopic is false and topic is not nil - set the topic to the value in topic - // Case 4: if ResetTopic is false and topic is nil, then leave the topic as is log.Info("generate segment updates without id", zap.Any("in", in)) ret := map[string]interface{}{} - if in.ResetTopic { - if in.Topic == nil { - ret["topic"] = nil - } - } else { - if in.Topic != nil { - ret["topic"] = *in.Topic - } - } - if in.ResetCollection { - if in.Collection == nil { - ret["collection_id"] = nil - } - } else { - if in.Collection != nil { - ret["collection_id"] = *in.Collection - } - } - log.Info("generate segment updates without id", zap.Any("updates", ret)) + // TODO: check this + //if in.ResetCollection { + // if in.Collection == nil { + // ret["collection_id"] = nil + // } + //} else { + // if in.Collection != nil { + // ret["collection_id"] = *in.Collection + // } + //} + //log.Info("generate segment updates without id", zap.Any("updates", ret)) return ret } func (s *segmentDb) Update(in *dbmodel.UpdateSegment) error { updates := generateSegmentUpdatesWithoutID(in) - return s.db.Model(&dbmodel.Segment{}).Where("id = ?", in.ID).Updates(updates).Error + return s.db.Model(&dbmodel.Segment{}). + Where("collection_id = ?", &in.Collection). + Where("id = ?", in.ID).Updates(updates).Error +} + +func (s *segmentDb) RegisterFilePaths(flushSegmentCompactions []*model.FlushSegmentCompaction) error { + log.Info("register file paths", zap.Any("flushSegmentCompactions", flushSegmentCompactions)) + for _, flushSegmentCompaction := range flushSegmentCompactions { + filePaths, err := json.Marshal(flushSegmentCompaction.FilePaths) + if err != nil { + log.Error("marshal file paths failed", zap.Error(err)) + return err + } + err = s.db.Model(&dbmodel.Segment{}). + Where("id = ?", flushSegmentCompaction.ID). + Update("file_paths", filePaths).Error + if err != nil { + log.Error("register file path failed", zap.Error(err)) + return err + } + } + return nil } diff --git a/go/coordinator/internal/metastore/db/dao/segment_metadata.go b/go/pkg/metastore/db/dao/segment_metadata.go similarity index 90% rename from go/coordinator/internal/metastore/db/dao/segment_metadata.go rename to go/pkg/metastore/db/dao/segment_metadata.go index 14d4d2ec2d0..cd4ace0efd8 100644 --- a/go/coordinator/internal/metastore/db/dao/segment_metadata.go +++ b/go/pkg/metastore/db/dao/segment_metadata.go @@ -1,7 +1,7 @@ package dao import ( - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" "gorm.io/gorm" "gorm.io/gorm/clause" ) @@ -21,7 +21,7 @@ func (s *segmentMetadataDb) DeleteBySegmentID(segmentID string) error { func (s *segmentMetadataDb) DeleteBySegmentIDAndKeys(segmentID string, keys []string) error { return s.db. Where("segment_id = ?", segmentID). - Where("`key` IN ?", keys). + Where("key IN ?", keys). Delete(&dbmodel.SegmentMetadata{}).Error } diff --git a/go/pkg/metastore/db/dao/segment_test.go b/go/pkg/metastore/db/dao/segment_test.go new file mode 100644 index 00000000000..f2a5cc0409c --- /dev/null +++ b/go/pkg/metastore/db/dao/segment_test.go @@ -0,0 +1,150 @@ +package dao + +import ( + "strconv" + "testing" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "k8s.io/apimachinery/pkg/util/rand" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" + "gorm.io/gorm" +) + +type SegmentDbTestSuite struct { + suite.Suite + db *gorm.DB + segmentDb *segmentDb +} + +func (suite *SegmentDbTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() + suite.segmentDb = &segmentDb{ + db: suite.db, + } +} + +func (suite *SegmentDbTestSuite) TestSegmentDb_GetSegments() { + uniqueID := types.NewUniqueID() + collectionID := uniqueID.String() + segment := &dbmodel.Segment{ + ID: uniqueID.String(), + CollectionID: &collectionID, + Type: "test_type", + Scope: "test_scope", + } + err := suite.db.Create(segment).Error + suite.NoError(err) + + testKey := "test" + testValue := "test" + metadata := &dbmodel.SegmentMetadata{ + SegmentID: segment.ID, + Key: &testKey, + StrValue: &testValue, + } + err = suite.db.Create(metadata).Error + suite.NoError(err) + + // Test when all parameters are nil + segments, err := suite.segmentDb.GetSegments(types.NilUniqueID(), nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Len(segments, 1) + suite.Equal(segment.ID, segments[0].Segment.ID) + suite.Equal(segment.CollectionID, segments[0].Segment.CollectionID) + suite.Equal(segment.Type, segments[0].Segment.Type) + suite.Equal(segment.Scope, segments[0].Segment.Scope) + suite.Len(segments[0].SegmentMetadata, 1) + suite.Equal(metadata.Key, segments[0].SegmentMetadata[0].Key) + suite.Equal(metadata.StrValue, segments[0].SegmentMetadata[0].StrValue) + + // Test when filtering by ID + segments, err = suite.segmentDb.GetSegments(types.MustParse(segment.ID), nil, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Len(segments, 1) + suite.Equal(segment.ID, segments[0].Segment.ID) + + // Test when filtering by type + segments, err = suite.segmentDb.GetSegments(types.NilUniqueID(), &segment.Type, nil, types.NilUniqueID()) + suite.NoError(err) + suite.Len(segments, 1) + suite.Equal(segment.ID, segments[0].Segment.ID) + + // Test when filtering by scope + segments, err = suite.segmentDb.GetSegments(types.NilUniqueID(), nil, &segment.Scope, types.NilUniqueID()) + suite.NoError(err) + suite.Len(segments, 1) + suite.Equal(segment.ID, segments[0].Segment.ID) + + // Test when filtering by collection ID + segments, err = suite.segmentDb.GetSegments(types.NilUniqueID(), nil, nil, types.MustParse(*segment.CollectionID)) + suite.NoError(err) + suite.Len(segments, 1) + suite.Equal(segment.ID, segments[0].Segment.ID) + + // clean up + err = suite.db.Delete(segment).Error + suite.NoError(err) + err = suite.db.Delete(metadata).Error + suite.NoError(err) +} + +func (suite *SegmentDbTestSuite) TestSegmentDb_RegisterFilePath() { + // create a collection for testing + databaseId := types.NewUniqueID().String() + collectionName := "test_segment_register_file_paths" + collectionID, err := CreateTestCollection(suite.db, collectionName, 128, databaseId) + suite.NoError(err) + + segments, err := suite.segmentDb.GetSegments(types.NilUniqueID(), nil, nil, types.MustParse(collectionID)) + suite.NoError(err) + + // create entries to flush + segmentsFilePaths := make(map[string]map[string][]string) + flushSegmentCompactions := make([]*model.FlushSegmentCompaction, 0) + testFilePathTypes := []string{"TypeA", "TypeB", "TypeC", "TypeD"} + for _, segment := range segments { + segmentID := segment.Segment.ID + segmentsFilePaths[segmentID] = make(map[string][]string) + for i := 0; i < rand.Intn(len(testFilePathTypes)); i++ { + filePaths := make([]string, 0) + for j := 0; j < rand.Intn(5); j++ { + filePaths = append(filePaths, "test_file_path_"+strconv.Itoa(j+1)) + } + filePathTypeI := rand.Intn(len(testFilePathTypes)) + filePathType := testFilePathTypes[filePathTypeI] + segmentsFilePaths[segmentID][filePathType] = filePaths + } + flushSegmentCompaction := &model.FlushSegmentCompaction{ + ID: types.MustParse(segmentID), + FilePaths: segmentsFilePaths[segmentID], + } + flushSegmentCompactions = append(flushSegmentCompactions, flushSegmentCompaction) + } + + // flush the entries + err = suite.segmentDb.RegisterFilePaths(flushSegmentCompactions) + suite.NoError(err) + + // verify file paths registered + segments, err = suite.segmentDb.GetSegments(types.NilUniqueID(), nil, nil, types.MustParse(collectionID)) + suite.NoError(err) + for _, segment := range segments { + suite.Contains(segmentsFilePaths, segment.Segment.ID) + suite.Equal(segmentsFilePaths[segment.Segment.ID], segment.Segment.FilePaths) + } + + // clean up + err = CleanUpTestCollection(suite.db, collectionID) + suite.NoError(err) +} + +func TestSegmentDbTestSuiteSuite(t *testing.T) { + testSuite := new(SegmentDbTestSuite) + suite.Run(t, testSuite) +} diff --git a/go/pkg/metastore/db/dao/tenant.go b/go/pkg/metastore/db/dao/tenant.go new file mode 100644 index 00000000000..fcd73f2cdcb --- /dev/null +++ b/go/pkg/metastore/db/dao/tenant.go @@ -0,0 +1,98 @@ +package dao + +import ( + "errors" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/jackc/pgx/v5/pgconn" + "github.com/pingcap/log" + "go.uber.org/zap" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +type tenantDb struct { + db *gorm.DB +} + +var _ dbmodel.ITenantDb = &tenantDb{} + +func (s *tenantDb) DeleteAll() error { + return s.db.Where("1 = 1").Delete(&dbmodel.Tenant{}).Error +} + +func (s *tenantDb) DeleteByID(tenantID string) (int, error) { + var tenants []dbmodel.Tenant + err := s.db.Clauses(clause.Returning{}).Where("id = ?", tenantID).Delete(&tenants).Error + return len(tenants), err +} + +func (s *tenantDb) GetAllTenants() ([]*dbmodel.Tenant, error) { + var tenants []*dbmodel.Tenant + + if err := s.db.Find(&tenants).Error; err != nil { + return nil, err + } + return tenants, nil +} + +func (s *tenantDb) GetTenants(tenantID string) ([]*dbmodel.Tenant, error) { + var tenants []*dbmodel.Tenant + + if err := s.db.Where("id = ?", tenantID).Find(&tenants).Error; err != nil { + return nil, err + } + return tenants, nil +} + +func (s *tenantDb) Insert(tenant *dbmodel.Tenant) error { + err := s.db.Create(tenant).Error + if err != nil { + log.Error("create tenant failed", zap.Error(err)) + var pgErr *pgconn.PgError + ok := errors.As(err, &pgErr) + if ok { + log.Error("Postgres Error") + switch pgErr.Code { + case "23505": + log.Error("tenant already exists") + return common.ErrTenantUniqueConstraintViolation + default: + return err + } + } + return err + } + return nil +} + +func (s *tenantDb) UpdateTenantLastCompactionTime(tenantID string, lastCompactionTime int64) error { + log.Info("UpdateTenantLastCompactionTime", zap.String("tenantID", tenantID), zap.Int64("lastCompactionTime", lastCompactionTime)) + var tenants []dbmodel.Tenant + result := s.db.Model(&tenants). + Clauses(clause.Returning{Columns: []clause.Column{{Name: "id"}}}). + Where("id = ?", tenantID). + Update("last_compaction_time", lastCompactionTime) + + if result.Error != nil { + log.Error("UpdateTenantLastCompactionTime error", zap.Error(result.Error)) + return result.Error + } + if result.RowsAffected == 0 { + return common.ErrTenantNotFound + } + return nil +} + +func (s *tenantDb) GetTenantsLastCompactionTime(tenantIDs []string) ([]*dbmodel.Tenant, error) { + log.Info("GetTenantsLastCompactionTime", zap.Any("tenantIDs", tenantIDs)) + var tenants []*dbmodel.Tenant + + result := s.db.Select("id", "last_compaction_time").Find(&tenants, "id IN ?", tenantIDs) + if result.Error != nil { + log.Error("GetTenantsLastCompactionTime error", zap.Error(result.Error)) + return nil, result.Error + } + + return tenants, nil +} diff --git a/go/pkg/metastore/db/dao/tenant_test.go b/go/pkg/metastore/db/dao/tenant_test.go new file mode 100644 index 00000000000..7bb613ae7df --- /dev/null +++ b/go/pkg/metastore/db/dao/tenant_test.go @@ -0,0 +1,102 @@ +package dao + +import ( + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/pingcap/log" + "github.com/stretchr/testify/suite" + "gorm.io/gorm" + "strconv" + "testing" + "time" +) + +type TenantDbTestSuite struct { + suite.Suite + db *gorm.DB + Db *tenantDb + t *testing.T +} + +func (suite *TenantDbTestSuite) SetupSuite() { + log.Info("setup suite") + suite.db = dbcore.ConfigDatabaseForTesting() + suite.Db = &tenantDb{ + db: suite.db, + } +} + +func (suite *TenantDbTestSuite) SetupTest() { + log.Info("setup test") +} + +func (suite *TenantDbTestSuite) TearDownTest() { + log.Info("teardown test") +} + +func (suite *TenantDbTestSuite) TestTenantDb_UpdateTenantLastCompactionTime() { + tenantId := "testUpdateTenantLastCompactionTime" + var tenant dbmodel.Tenant + err := suite.Db.Insert(&dbmodel.Tenant{ + ID: tenantId, + LastCompactionTime: 0, + }) + suite.Require().NoError(err) + suite.db.First(&tenant, "id = ?", tenantId) + suite.Require().Equal(int64(0), tenant.LastCompactionTime) + + err = suite.Db.UpdateTenantLastCompactionTime(tenantId, 1) + suite.Require().NoError(err) + suite.db.First(&tenant, "id = ?", tenantId) + suite.Require().Equal(int64(1), tenant.LastCompactionTime) + + currentTime := time.Now().Unix() + err = suite.Db.UpdateTenantLastCompactionTime(tenantId, currentTime) + suite.Require().NoError(err) + suite.db.First(&tenant, "id = ?", tenantId) + suite.Require().Equal(currentTime, tenant.LastCompactionTime) + + suite.db.Delete(&tenant, "id = ?", tenantId) +} + +func (suite *TenantDbTestSuite) TestTenantDb_GetTenantsLastCompactionTime() { + tenantIds := make([]string, 0) + for i := 0; i < 10; i++ { + tenantId := "testGetTenantsLastCompactionTime" + strconv.Itoa(i) + err := suite.Db.Insert(&dbmodel.Tenant{ + ID: tenantId, + LastCompactionTime: int64(i), + }) + suite.Require().NoError(err) + tenantIds = append(tenantIds, tenantId) + } + + tenants, err := suite.Db.GetTenantsLastCompactionTime(tenantIds) + suite.Require().NoError(err) + suite.Require().Len(tenants, 10) + for i, tenant := range tenants { + suite.Require().Equal(int64(i), tenant.LastCompactionTime) + } + + currentTime := time.Now().Unix() + for _, tenantId := range tenantIds { + err := suite.Db.UpdateTenantLastCompactionTime(tenantId, currentTime) + suite.Require().NoError(err) + } + tenants, err = suite.Db.GetTenantsLastCompactionTime(tenantIds) + suite.Require().NoError(err) + suite.Require().Len(tenants, 10) + for _, tenant := range tenants { + suite.Require().Equal(currentTime, tenant.LastCompactionTime) + } + + for _, tenantId := range tenantIds { + suite.db.Delete(&dbmodel.Tenant{}, "id = ?", tenantId) + } +} + +func TestTenantDbTestSuite(t *testing.T) { + testSuite := new(TenantDbTestSuite) + testSuite.t = t + suite.Run(t, testSuite) +} diff --git a/go/pkg/metastore/db/dao/test_utils.go b/go/pkg/metastore/db/dao/test_utils.go new file mode 100644 index 00000000000..874dcabc112 --- /dev/null +++ b/go/pkg/metastore/db/dao/test_utils.go @@ -0,0 +1,184 @@ +package dao + +import ( + "time" + + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/types" + "github.com/pingcap/log" + "go.uber.org/zap" + "gorm.io/gorm" +) + +const SegmentType = "urn:chroma:segment/vector/hnsw-distributed" + +func GetSegmentScopes() []string { + return []string{"VECTOR", "METADATA"} +} + +func CreateTestTenantAndDatabase(db *gorm.DB, tenant string, database string) (string, error) { + log.Info("create test tenant and database", zap.String("tenant", tenant), zap.String("database", database)) + tenantDb := &tenantDb{ + db: db, + } + databaseDb := &databaseDb{ + db: db, + } + + err := tenantDb.Insert(&dbmodel.Tenant{ + ID: tenant, + LastCompactionTime: time.Now().Unix(), + }) + if err != nil { + return "", err + } + + databaseId := types.NewUniqueID().String() + err = databaseDb.Insert(&dbmodel.Database{ + ID: databaseId, + Name: database, + TenantID: tenant, + }) + if err != nil { + return "", err + } + + return databaseId, nil +} + +func CleanUpTestDatabase(db *gorm.DB, tenantName string, databaseName string) error { + log.Info("clean up test database", zap.String("tenantName", tenantName), zap.String("databaseName", databaseName)) + // clean up collections + collectionDb := &collectionDb{ + db: db, + } + collections, err := collectionDb.GetCollections(nil, nil, tenantName, databaseName) + log.Info("clean up test database", zap.Int("collections", len(collections))) + if err != nil { + return err + } + for _, collection := range collections { + err = CleanUpTestCollection(db, collection.Collection.ID) + if err != nil { + return err + } + } + + // clean up database + databaseDb := &databaseDb{ + db: db, + } + + _, err = databaseDb.DeleteByTenantIdAndName(tenantName, databaseName) + if err != nil { + return err + } + + return nil +} + +func CleanUpTestTenant(db *gorm.DB, tenantName string) error { + log.Info("clean up test tenant", zap.String("tenantName", tenantName)) + tenantDb := &tenantDb{ + db: db, + } + databaseDb := &databaseDb{ + db: db, + } + + // clean up databases + databases, err := databaseDb.GetDatabasesByTenantID(tenantName) + if err != nil { + return err + } + for _, database := range databases { + err = CleanUpTestDatabase(db, tenantName, database.Name) + if err != nil { + return err + } + } + + // clean up tenant + _, err = tenantDb.DeleteByID(tenantName) + if err != nil { + return err + } + return nil +} + +func CreateTestCollection(db *gorm.DB, collectionName string, dimension int32, databaseID string) (string, error) { + log.Info("create test collection", zap.String("collectionName", collectionName), zap.Int32("dimension", dimension), zap.String("databaseID", databaseID)) + collectionDb := &collectionDb{ + db: db, + } + segmentDb := &segmentDb{ + db: db, + } + collectionId := types.NewUniqueID().String() + + err := collectionDb.Insert(&dbmodel.Collection{ + ID: collectionId, + Name: &collectionName, + Dimension: &dimension, + DatabaseID: databaseID, + }) + if err != nil { + return "", err + } + + for _, scope := range GetSegmentScopes() { + segmentId := types.NewUniqueID().String() + err = segmentDb.Insert(&dbmodel.Segment{ + CollectionID: &collectionId, + ID: segmentId, + Type: SegmentType, + Scope: scope, + }) + if err != nil { + return "", err + } + } + + return collectionId, nil +} + +func CleanUpTestCollection(db *gorm.DB, collectionId string) error { + log.Info("clean up collection", zap.String("collectionId", collectionId)) + collectionDb := &collectionDb{ + db: db, + } + collectionMetadataDb := &collectionMetadataDb{ + db: db, + } + segmentDb := &segmentDb{ + db: db, + } + segmentMetadataDb := &segmentMetadataDb{ + db: db, + } + + _, err := collectionMetadataDb.DeleteByCollectionID(collectionId) + if err != nil { + return err + } + _, err = collectionDb.DeleteCollectionByID(collectionId) + if err != nil { + return err + } + segments, err := segmentDb.GetSegments(types.NilUniqueID(), nil, nil, types.MustParse(collectionId)) + if err != nil { + return err + } + for _, segment := range segments { + err = segmentDb.DeleteSegmentByID(segment.Segment.ID) + if err != nil { + return err + } + err = segmentMetadataDb.DeleteBySegmentID(segment.Segment.ID) + if err != nil { + return err + } + } + + return nil +} diff --git a/go/pkg/metastore/db/dbcore/core.go b/go/pkg/metastore/db/dbcore/core.go new file mode 100644 index 00000000000..956df311253 --- /dev/null +++ b/go/pkg/metastore/db/dbcore/core.go @@ -0,0 +1,213 @@ +package dbcore + +import ( + "context" + "fmt" + "os" + "reflect" + "strconv" + "time" + + "github.com/chroma-core/chroma/go/pkg/types" + + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/pingcap/log" + "go.uber.org/zap" + "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +var ( + globalDB *gorm.DB +) + +type DBConfig struct { + Username string + Password string + Address string + Port int + DBName string + MaxIdleConns int + MaxOpenConns int + SslMode string +} + +func ConnectPostgres(cfg DBConfig) (*gorm.DB, error) { + log.Info("ConnectPostgres", zap.String("host", cfg.Address), zap.String("database", cfg.DBName), zap.Int("port", cfg.Port)) + dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d sslmode=%s", + cfg.Address, cfg.Username, cfg.Password, cfg.DBName, cfg.Port, cfg.SslMode) + + ormLogger := logger.Default + ormLogger.LogMode(logger.Info) + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{ + Logger: ormLogger, + CreateBatchSize: 100, + }) + if err != nil { + log.Error("fail to connect db", + zap.String("host", cfg.Address), + zap.String("database", cfg.DBName), + zap.Error(err)) + return nil, err + } + + idb, err := db.DB() + if err != nil { + log.Error("fail to create db instance", + zap.String("host", cfg.Address), + zap.String("database", cfg.DBName), + zap.Error(err)) + return nil, err + } + idb.SetMaxIdleConns(cfg.MaxIdleConns) + idb.SetMaxOpenConns(cfg.MaxOpenConns) + + globalDB = db + + log.Info("Postgres connected success", + zap.String("host", cfg.Address), + zap.String("database", cfg.DBName), + zap.Error(err)) + + return db, nil +} + +// SetGlobalDB Only for test +func SetGlobalDB(db *gorm.DB) { + globalDB = db +} + +type ctxTransactionKey struct{} + +func CtxWithTransaction(ctx context.Context, tx *gorm.DB) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, ctxTransactionKey{}, tx) +} + +type txImpl struct{} + +func NewTxImpl() *txImpl { + return &txImpl{} +} + +func (*txImpl) Transaction(ctx context.Context, fn func(txctx context.Context) error) error { + db := globalDB.WithContext(ctx) + + return db.Transaction(func(tx *gorm.DB) error { + txCtx := CtxWithTransaction(ctx, tx) + return fn(txCtx) + }) +} + +func GetDB(ctx context.Context) *gorm.DB { + iface := ctx.Value(ctxTransactionKey{}) + + if iface != nil { + tx, ok := iface.(*gorm.DB) + if !ok { + log.Error("unexpect context value type", zap.Any("type", reflect.TypeOf(tx))) + return nil + } + + return tx + } + + return globalDB.WithContext(ctx) +} + +func CreateDefaultTenantAndDatabase(db *gorm.DB) string { + defaultTenant := &dbmodel.Tenant{ + ID: common.DefaultTenant, + LastCompactionTime: time.Now().Unix(), + } + db.Model(&dbmodel.Tenant{}).Where("id = ?", common.DefaultTenant).Save(defaultTenant) + + var database []dbmodel.Database + databaseId := types.NewUniqueID().String() + result := db.Model(&dbmodel.Database{}). + Where("name = ?", common.DefaultDatabase). + Where("tenant_id = ?", common.DefaultTenant). + Find(&database) + if result.Error != nil { + return "" + } + + if result.RowsAffected == 0 { + db.Create(&dbmodel.Database{ + ID: databaseId, + Name: common.DefaultDatabase, + TenantID: common.DefaultTenant, + }) + return databaseId + } + + err := result.Row().Scan(&database) + if err != nil { + return "" + } + return database[0].ID +} + +func CreateTestTables(db *gorm.DB) { + log.Info("CreateTestTables") + tableExist := db.Migrator().HasTable(&dbmodel.Tenant{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.Tenant{}) + } + tableExist = db.Migrator().HasTable(&dbmodel.Database{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.Database{}) + } + tableExist = db.Migrator().HasTable(&dbmodel.CollectionMetadata{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.CollectionMetadata{}) + } + tableExist = db.Migrator().HasTable(&dbmodel.Collection{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.Collection{}) + } + tableExist = db.Migrator().HasTable(&dbmodel.SegmentMetadata{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.SegmentMetadata{}) + } + tableExist = db.Migrator().HasTable(&dbmodel.Segment{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.Segment{}) + } + tableExist = db.Migrator().HasTable(&dbmodel.Notification{}) + if !tableExist { + db.Migrator().CreateTable(&dbmodel.Notification{}) + } + + // create default tenant and database + CreateDefaultTenantAndDatabase(db) +} + +func GetDBConfigForTesting() DBConfig { + dbAddress := os.Getenv("POSTGRES_HOST") + dbPort, _ := strconv.Atoi(os.Getenv("POSTGRES_PORT")) + return DBConfig{ + Username: "chroma", + Password: "chroma", + Address: dbAddress, + Port: dbPort, + DBName: "chroma", + MaxIdleConns: 10, + MaxOpenConns: 100, + SslMode: "disable", + } +} + +func ConfigDatabaseForTesting() *gorm.DB { + db, err := ConnectPostgres(GetDBConfigForTesting()) + if err != nil { + panic("failed to connect database") + } + SetGlobalDB(db) + CreateTestTables(db) + return db +} diff --git a/go/pkg/metastore/db/dbmodel/collection.go b/go/pkg/metastore/db/dbmodel/collection.go new file mode 100644 index 00000000000..c6c769b4fa2 --- /dev/null +++ b/go/pkg/metastore/db/dbmodel/collection.go @@ -0,0 +1,41 @@ +package dbmodel + +import ( + "time" + + "github.com/chroma-core/chroma/go/pkg/types" +) + +type Collection struct { + ID string `gorm:"id;primaryKey"` + Name *string `gorm:"name;unique"` + Dimension *int32 `gorm:"dimension"` + DatabaseID string `gorm:"database_id"` + Ts types.Timestamp `gorm:"ts;type:bigint;default:0"` + IsDeleted bool `gorm:"is_deleted;type:bool;default:false"` + CreatedAt time.Time `gorm:"created_at;type:timestamp;not null;default:current_timestamp"` + UpdatedAt time.Time `gorm:"updated_at;type:timestamp;not null;default:current_timestamp"` + LogPosition int64 `gorm:"log_position;default:0"` + Version int32 `gorm:"version;default:0"` +} + +func (v Collection) TableName() string { + return "collections" +} + +type CollectionAndMetadata struct { + Collection *Collection + CollectionMetadata []*CollectionMetadata + TenantID string + DatabaseName string +} + +//go:generate mockery --name=ICollectionDb +type ICollectionDb interface { + GetCollections(collectionID *string, collectionName *string, tenantID string, databaseName string) ([]*CollectionAndMetadata, error) + DeleteCollectionByID(collectionID string) (int, error) + Insert(in *Collection) error + Update(in *Collection) error + DeleteAll() error + UpdateLogPositionAndVersion(collectionID string, logPosition int64, currentCollectionVersion int32) (int32, error) +} diff --git a/go/coordinator/internal/metastore/db/dbmodel/collection_metadata.go b/go/pkg/metastore/db/dbmodel/collection_metadata.go similarity index 88% rename from go/coordinator/internal/metastore/db/dbmodel/collection_metadata.go rename to go/pkg/metastore/db/dbmodel/collection_metadata.go index 29303453c5d..6c519a731d1 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/collection_metadata.go +++ b/go/pkg/metastore/db/dbmodel/collection_metadata.go @@ -3,7 +3,7 @@ package dbmodel import ( "time" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/types" ) type CollectionMetadata struct { @@ -23,7 +23,7 @@ func (v CollectionMetadata) TableName() string { //go:generate mockery --name=ICollectionMetadataDb type ICollectionMetadataDb interface { - DeleteByCollectionID(collectionID string) error + DeleteByCollectionID(collectionID string) (int, error) Insert(in []*CollectionMetadata) error DeleteAll() error } diff --git a/go/coordinator/internal/metastore/db/dbmodel/common.go b/go/pkg/metastore/db/dbmodel/common.go similarity index 93% rename from go/coordinator/internal/metastore/db/dbmodel/common.go rename to go/pkg/metastore/db/dbmodel/common.go index d188193ae18..d90b7df55e6 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/common.go +++ b/go/pkg/metastore/db/dbmodel/common.go @@ -15,6 +15,7 @@ type IMetaDomain interface { SegmentDb(ctx context.Context) ISegmentDb SegmentMetadataDb(ctx context.Context) ISegmentMetadataDb NotificationDb(ctx context.Context) INotificationDb + RecordLogDb(ctx context.Context) IRecordLogDb } //go:generate mockery --name=ITransaction diff --git a/go/coordinator/internal/metastore/db/dbmodel/database.go b/go/pkg/metastore/db/dbmodel/database.go similarity index 94% rename from go/coordinator/internal/metastore/db/dbmodel/database.go rename to go/pkg/metastore/db/dbmodel/database.go index 6cac848b423..9413387ec0e 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/database.go +++ b/go/pkg/metastore/db/dbmodel/database.go @@ -3,7 +3,7 @@ package dbmodel import ( "time" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/types" ) type Database struct { diff --git a/go/pkg/metastore/db/dbmodel/mocks/ICollectionDb.go b/go/pkg/metastore/db/dbmodel/mocks/ICollectionDb.go new file mode 100644 index 00000000000..18624756268 --- /dev/null +++ b/go/pkg/metastore/db/dbmodel/mocks/ICollectionDb.go @@ -0,0 +1,167 @@ +// Code generated by mockery v2.42.0. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// ICollectionDb is an autogenerated mock type for the ICollectionDb type +type ICollectionDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *ICollectionDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteCollectionByID provides a mock function with given fields: collectionID +func (_m *ICollectionDb) DeleteCollectionByID(collectionID string) (int, error) { + ret := _m.Called(collectionID) + + if len(ret) == 0 { + panic("no return value specified for DeleteCollectionByID") + } + + var r0 int + var r1 error + if rf, ok := ret.Get(0).(func(string) (int, error)); ok { + return rf(collectionID) + } + if rf, ok := ret.Get(0).(func(string) int); ok { + r0 = rf(collectionID) + } else { + r0 = ret.Get(0).(int) + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetCollections provides a mock function with given fields: collectionID, collectionName, tenantID, databaseName +func (_m *ICollectionDb) GetCollections(collectionID *string, collectionName *string, tenantID string, databaseName string) ([]*dbmodel.CollectionAndMetadata, error) { + ret := _m.Called(collectionID, collectionName, tenantID, databaseName) + + if len(ret) == 0 { + panic("no return value specified for GetCollections") + } + + var r0 []*dbmodel.CollectionAndMetadata + var r1 error + if rf, ok := ret.Get(0).(func(*string, *string, string, string) ([]*dbmodel.CollectionAndMetadata, error)); ok { + return rf(collectionID, collectionName, tenantID, databaseName) + } + if rf, ok := ret.Get(0).(func(*string, *string, string, string) []*dbmodel.CollectionAndMetadata); ok { + r0 = rf(collectionID, collectionName, tenantID, databaseName) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*dbmodel.CollectionAndMetadata) + } + } + + if rf, ok := ret.Get(1).(func(*string, *string, string, string) error); ok { + r1 = rf(collectionID, collectionName, tenantID, databaseName) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: in +func (_m *ICollectionDb) Insert(in *dbmodel.Collection) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Collection) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Update provides a mock function with given fields: in +func (_m *ICollectionDb) Update(in *dbmodel.Collection) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Update") + } + + var r0 error + if rf, ok := ret.Get(0).(func(*dbmodel.Collection) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// UpdateLogPositionAndVersion provides a mock function with given fields: collectionID, logPosition, currentCollectionVersion +func (_m *ICollectionDb) UpdateLogPositionAndVersion(collectionID string, logPosition int64, currentCollectionVersion int32) (int32, error) { + ret := _m.Called(collectionID, logPosition, currentCollectionVersion) + + if len(ret) == 0 { + panic("no return value specified for UpdateLogPositionAndVersion") + } + + var r0 int32 + var r1 error + if rf, ok := ret.Get(0).(func(string, int64, int32) (int32, error)); ok { + return rf(collectionID, logPosition, currentCollectionVersion) + } + if rf, ok := ret.Get(0).(func(string, int64, int32) int32); ok { + r0 = rf(collectionID, logPosition, currentCollectionVersion) + } else { + r0 = ret.Get(0).(int32) + } + + if rf, ok := ret.Get(1).(func(string, int64, int32) error); ok { + r1 = rf(collectionID, logPosition, currentCollectionVersion) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewICollectionDb creates a new instance of ICollectionDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewICollectionDb(t interface { + mock.TestingT + Cleanup(func()) +}) *ICollectionDb { + mock := &ICollectionDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/pkg/metastore/db/dbmodel/mocks/ICollectionMetadataDb.go b/go/pkg/metastore/db/dbmodel/mocks/ICollectionMetadataDb.go new file mode 100644 index 00000000000..d89e09536d9 --- /dev/null +++ b/go/pkg/metastore/db/dbmodel/mocks/ICollectionMetadataDb.go @@ -0,0 +1,91 @@ +// Code generated by mockery v2.42.0. DO NOT EDIT. + +package mocks + +import ( + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + mock "github.com/stretchr/testify/mock" +) + +// ICollectionMetadataDb is an autogenerated mock type for the ICollectionMetadataDb type +type ICollectionMetadataDb struct { + mock.Mock +} + +// DeleteAll provides a mock function with given fields: +func (_m *ICollectionMetadataDb) DeleteAll() error { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for DeleteAll") + } + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// DeleteByCollectionID provides a mock function with given fields: collectionID +func (_m *ICollectionMetadataDb) DeleteByCollectionID(collectionID string) (int, error) { + ret := _m.Called(collectionID) + + if len(ret) == 0 { + panic("no return value specified for DeleteByCollectionID") + } + + var r0 int + var r1 error + if rf, ok := ret.Get(0).(func(string) (int, error)); ok { + return rf(collectionID) + } + if rf, ok := ret.Get(0).(func(string) int); ok { + r0 = rf(collectionID) + } else { + r0 = ret.Get(0).(int) + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(collectionID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Insert provides a mock function with given fields: in +func (_m *ICollectionMetadataDb) Insert(in []*dbmodel.CollectionMetadata) error { + ret := _m.Called(in) + + if len(ret) == 0 { + panic("no return value specified for Insert") + } + + var r0 error + if rf, ok := ret.Get(0).(func([]*dbmodel.CollectionMetadata) error); ok { + r0 = rf(in) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewICollectionMetadataDb creates a new instance of ICollectionMetadataDb. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewICollectionMetadataDb(t interface { + mock.TestingT + Cleanup(func()) +}) *ICollectionMetadataDb { + mock := &ICollectionMetadataDb{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/IDatabaseDb.go b/go/pkg/metastore/db/dbmodel/mocks/IDatabaseDb.go similarity index 96% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/IDatabaseDb.go rename to go/pkg/metastore/db/dbmodel/mocks/IDatabaseDb.go index 4bb8c5fa50c..712ec27c1cc 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/IDatabaseDb.go +++ b/go/pkg/metastore/db/dbmodel/mocks/IDatabaseDb.go @@ -3,7 +3,7 @@ package mocks import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" ) diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/IMetaDomain.go b/go/pkg/metastore/db/dbmodel/mocks/IMetaDomain.go similarity index 89% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/IMetaDomain.go rename to go/pkg/metastore/db/dbmodel/mocks/IMetaDomain.go index 0ee94c373e9..81803b79ebe 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/IMetaDomain.go +++ b/go/pkg/metastore/db/dbmodel/mocks/IMetaDomain.go @@ -5,7 +5,7 @@ package mocks import ( context "context" - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" ) @@ -126,6 +126,21 @@ func (_m *IMetaDomain) TenantDb(ctx context.Context) dbmodel.ITenantDb { return r0 } +func (_m *IMetaDomain) RecordLogDb(ctx context.Context) dbmodel.IRecordLogDb { + ret := _m.Called(ctx) + + var r0 dbmodel.IRecordLogDb + if rf, ok := ret.Get(0).(func(context.Context) dbmodel.IRecordLogDb); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(dbmodel.IRecordLogDb) + } + } + + return r0 +} + // NewIMetaDomain creates a new instance of IMetaDomain. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. // The first argument is typically a *testing.T value. func NewIMetaDomain(t interface { diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/INotificationDb.go b/go/pkg/metastore/db/dbmodel/mocks/INotificationDb.go similarity index 97% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/INotificationDb.go rename to go/pkg/metastore/db/dbmodel/mocks/INotificationDb.go index b5b9f77b394..b6a9cacdf38 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/INotificationDb.go +++ b/go/pkg/metastore/db/dbmodel/mocks/INotificationDb.go @@ -3,7 +3,7 @@ package mocks import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" ) diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/ISegmentDb.go b/go/pkg/metastore/db/dbmodel/mocks/ISegmentDb.go similarity index 76% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/ISegmentDb.go rename to go/pkg/metastore/db/dbmodel/mocks/ISegmentDb.go index 1a519766bba..5fa22500ce3 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/ISegmentDb.go +++ b/go/pkg/metastore/db/dbmodel/mocks/ISegmentDb.go @@ -3,10 +3,10 @@ package mocks import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" - types "github.com/chroma/chroma-coordinator/internal/types" + types "github.com/chroma-core/chroma/go/pkg/types" ) // ISegmentDb is an autogenerated mock type for the ISegmentDb type @@ -42,25 +42,25 @@ func (_m *ISegmentDb) DeleteSegmentByID(id string) error { return r0 } -// GetSegments provides a mock function with given fields: id, segmentType, scope, topic, collectionID -func (_m *ISegmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error) { - ret := _m.Called(id, segmentType, scope, topic, collectionID) +// GetSegments provides a mock function with given fields: id, segmentType, scope, collectionID +func (_m *ISegmentDb) GetSegments(id types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error) { + ret := _m.Called(id, segmentType, scope, collectionID) var r0 []*dbmodel.SegmentAndMetadata var r1 error - if rf, ok := ret.Get(0).(func(types.UniqueID, *string, *string, *string, types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error)); ok { - return rf(id, segmentType, scope, topic, collectionID) + if rf, ok := ret.Get(0).(func(types.UniqueID, *string, *string, types.UniqueID) ([]*dbmodel.SegmentAndMetadata, error)); ok { + return rf(id, segmentType, scope, collectionID) } - if rf, ok := ret.Get(0).(func(types.UniqueID, *string, *string, *string, types.UniqueID) []*dbmodel.SegmentAndMetadata); ok { - r0 = rf(id, segmentType, scope, topic, collectionID) + if rf, ok := ret.Get(0).(func(types.UniqueID, *string, *string, types.UniqueID) []*dbmodel.SegmentAndMetadata); ok { + r0 = rf(id, segmentType, scope, collectionID) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]*dbmodel.SegmentAndMetadata) } } - if rf, ok := ret.Get(1).(func(types.UniqueID, *string, *string, *string, types.UniqueID) error); ok { - r1 = rf(id, segmentType, scope, topic, collectionID) + if rf, ok := ret.Get(1).(func(types.UniqueID, *string, *string, types.UniqueID) error); ok { + r1 = rf(id, segmentType, scope, collectionID) } else { r1 = ret.Error(1) } diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/ISegmentMetadataDb.go b/go/pkg/metastore/db/dbmodel/mocks/ISegmentMetadataDb.go similarity index 96% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/ISegmentMetadataDb.go rename to go/pkg/metastore/db/dbmodel/mocks/ISegmentMetadataDb.go index 24c56b6d835..81c6cd01ca9 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/ISegmentMetadataDb.go +++ b/go/pkg/metastore/db/dbmodel/mocks/ISegmentMetadataDb.go @@ -3,7 +3,7 @@ package mocks import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" ) diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/ITenantDb.go b/go/pkg/metastore/db/dbmodel/mocks/ITenantDb.go similarity index 96% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/ITenantDb.go rename to go/pkg/metastore/db/dbmodel/mocks/ITenantDb.go index fe54c815037..743c9d18c8a 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/mocks/ITenantDb.go +++ b/go/pkg/metastore/db/dbmodel/mocks/ITenantDb.go @@ -3,7 +3,7 @@ package mocks import ( - dbmodel "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" + dbmodel "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" mock "github.com/stretchr/testify/mock" ) diff --git a/go/coordinator/internal/metastore/db/dbmodel/mocks/ITransaction.go b/go/pkg/metastore/db/dbmodel/mocks/ITransaction.go similarity index 100% rename from go/coordinator/internal/metastore/db/dbmodel/mocks/ITransaction.go rename to go/pkg/metastore/db/dbmodel/mocks/ITransaction.go diff --git a/go/coordinator/internal/metastore/db/dbmodel/notification.go b/go/pkg/metastore/db/dbmodel/notification.go similarity index 100% rename from go/coordinator/internal/metastore/db/dbmodel/notification.go rename to go/pkg/metastore/db/dbmodel/notification.go diff --git a/go/pkg/metastore/db/dbmodel/record_log.go b/go/pkg/metastore/db/dbmodel/record_log.go new file mode 100644 index 00000000000..221235c0b4d --- /dev/null +++ b/go/pkg/metastore/db/dbmodel/record_log.go @@ -0,0 +1,23 @@ +package dbmodel + +import ( + "github.com/chroma-core/chroma/go/pkg/types" +) + +type RecordLog struct { + CollectionID *string `gorm:"collection_id;primaryKey;autoIncrement:false"` + LogOffset int64 `gorm:"log_offset;primaryKey;autoIncrement:false"` + Timestamp int64 `gorm:"timestamp;"` + Record *[]byte `gorm:"record;type:bytea"` +} + +func (v RecordLog) TableName() string { + return "record_logs" +} + +//go:generate mockery --name=IRecordLogDb +type IRecordLogDb interface { + PushLogs(collectionID types.UniqueID, recordsContent [][]byte) (int, error) + PullLogs(collectionID types.UniqueID, id int64, batchSize int, endTimestamp int64) ([]*RecordLog, error) + GetAllCollectionsToCompact() ([]*RecordLog, error) +} diff --git a/go/pkg/metastore/db/dbmodel/segment.go b/go/pkg/metastore/db/dbmodel/segment.go new file mode 100644 index 00000000000..e5a3af0a8b3 --- /dev/null +++ b/go/pkg/metastore/db/dbmodel/segment.go @@ -0,0 +1,50 @@ +package dbmodel + +import ( + "time" + + "github.com/chroma-core/chroma/go/pkg/model" + + "github.com/chroma-core/chroma/go/pkg/types" +) + +type Segment struct { + /* Making CollectionID the primary key allows fast search when we have CollectionID. + This requires us to push down CollectionID from the caller. We don't think there is + need to modify CollectionID in the near future. Each Segment should always have a + collection as a parent and cannot be modified. */ + CollectionID *string `gorm:"collection_id;primaryKey"` + ID string `gorm:"id;primaryKey"` + Type string `gorm:"type;type:string;not null"` + Scope string `gorm:"scope"` + Ts types.Timestamp `gorm:"ts;type:bigint;default:0"` + IsDeleted bool `gorm:"is_deleted;type:bool;default:false"` + CreatedAt time.Time `gorm:"created_at;type:timestamp;not null;default:current_timestamp"` + UpdatedAt time.Time `gorm:"updated_at;type:timestamp;not null;default:current_timestamp"` + FilePaths map[string][]string `gorm:"file_paths;serializer:json;default:'{}'"` +} + +func (s Segment) TableName() string { + return "segments" +} + +type SegmentAndMetadata struct { + Segment *Segment + SegmentMetadata []*SegmentMetadata +} + +type UpdateSegment struct { + ID string + Collection *string + ResetCollection bool +} + +//go:generate mockery --name=ISegmentDb +type ISegmentDb interface { + GetSegments(id types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID) ([]*SegmentAndMetadata, error) + DeleteSegmentByID(id string) error + Insert(*Segment) error + Update(*UpdateSegment) error + DeleteAll() error + RegisterFilePaths(flushSegmentCompactions []*model.FlushSegmentCompaction) error +} diff --git a/go/coordinator/internal/metastore/db/dbmodel/segment_metadata.go b/go/pkg/metastore/db/dbmodel/segment_metadata.go similarity index 94% rename from go/coordinator/internal/metastore/db/dbmodel/segment_metadata.go rename to go/pkg/metastore/db/dbmodel/segment_metadata.go index bbd11eaa39b..3d830225716 100644 --- a/go/coordinator/internal/metastore/db/dbmodel/segment_metadata.go +++ b/go/pkg/metastore/db/dbmodel/segment_metadata.go @@ -3,7 +3,7 @@ package dbmodel import ( "time" - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/types" ) type SegmentMetadata struct { diff --git a/go/pkg/metastore/db/dbmodel/tenant.go b/go/pkg/metastore/db/dbmodel/tenant.go new file mode 100644 index 00000000000..1db2a223f7b --- /dev/null +++ b/go/pkg/metastore/db/dbmodel/tenant.go @@ -0,0 +1,30 @@ +package dbmodel + +import ( + "time" + + "github.com/chroma-core/chroma/go/pkg/types" +) + +type Tenant struct { + ID string `gorm:"id;primaryKey;unique"` + Ts types.Timestamp `gorm:"ts;type:bigint;default:0"` + IsDeleted bool `gorm:"is_deleted;type:bool;default:false"` + CreatedAt time.Time `gorm:"created_at;type:timestamp;not null;default:current_timestamp"` + UpdatedAt time.Time `gorm:"updated_at;type:timestamp;not null;default:current_timestamp"` + LastCompactionTime int64 `gorm:"last_compaction_time;not null"` +} + +func (v Tenant) TableName() string { + return "tenants" +} + +//go:generate mockery --name=ITenantDb +type ITenantDb interface { + GetAllTenants() ([]*Tenant, error) + GetTenants(tenantID string) ([]*Tenant, error) + Insert(in *Tenant) error + DeleteAll() error + UpdateTenantLastCompactionTime(tenantID string, lastCompactionTime int64) error + GetTenantsLastCompactionTime(tenantIDs []string) ([]*Tenant, error) +} diff --git a/go/coordinator/internal/metastore/mocks/Catalog.go b/go/pkg/metastore/mocks/Catalog.go similarity index 79% rename from go/coordinator/internal/metastore/mocks/Catalog.go rename to go/pkg/metastore/mocks/Catalog.go index 5926bc768f0..596d2b6a33f 100644 --- a/go/coordinator/internal/metastore/mocks/Catalog.go +++ b/go/pkg/metastore/mocks/Catalog.go @@ -7,9 +7,9 @@ import ( mock "github.com/stretchr/testify/mock" - model "github.com/chroma/chroma-coordinator/internal/model" + model "github.com/chroma-core/chroma/go/pkg/model" - types "github.com/chroma/chroma-coordinator/internal/types" + types "github.com/chroma-core/chroma/go/pkg/types" ) // Catalog is an autogenerated mock type for the Catalog type @@ -97,25 +97,25 @@ func (_m *Catalog) DeleteSegment(ctx context.Context, segmentID types.UniqueID) return r0 } -// GetCollections provides a mock function with given fields: ctx, collectionID, collectionName, collectionTopic -func (_m *Catalog) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string) ([]*model.Collection, error) { - ret := _m.Called(ctx, collectionID, collectionName, collectionTopic) +// GetCollections provides a mock function with given fields: ctx, collectionID, collectionName +func (_m *Catalog) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string) ([]*model.Collection, error) { + ret := _m.Called(ctx, collectionID, collectionName) var r0 []*model.Collection var r1 error - if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string) ([]*model.Collection, error)); ok { - return rf(ctx, collectionID, collectionName, collectionTopic) + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string) ([]*model.Collection, error)); ok { + return rf(ctx, collectionID, collectionName) } - if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string) []*model.Collection); ok { - r0 = rf(ctx, collectionID, collectionName, collectionTopic) + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string) []*model.Collection); ok { + r0 = rf(ctx, collectionID, collectionName) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]*model.Collection) } } - if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, *string) error); ok { - r1 = rf(ctx, collectionID, collectionName, collectionTopic) + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string) error); ok { + r1 = rf(ctx, collectionID, collectionName) } else { r1 = ret.Error(1) } @@ -123,25 +123,25 @@ func (_m *Catalog) GetCollections(ctx context.Context, collectionID types.Unique return r0, r1 } -// GetSegments provides a mock function with given fields: ctx, segmentID, segmentType, scope, topic, collectionID, ts -func (_m *Catalog) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID, ts int64) ([]*model.Segment, error) { - ret := _m.Called(ctx, segmentID, segmentType, scope, topic, collectionID, ts) +// GetSegments provides a mock function with given fields: ctx, segmentID, segmentType, scope, collectionID, ts +func (_m *Catalog) GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, collectionID types.UniqueID, ts int64) ([]*model.Segment, error) { + ret := _m.Called(ctx, segmentID, segmentType, scope, collectionID, ts) var r0 []*model.Segment var r1 error - if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, *string, types.UniqueID, int64) ([]*model.Segment, error)); ok { - return rf(ctx, segmentID, segmentType, scope, topic, collectionID, ts) + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID, int64) ([]*model.Segment, error)); ok { + return rf(ctx, segmentID, segmentType, scope, collectionID, ts) } - if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, *string, types.UniqueID, int64) []*model.Segment); ok { - r0 = rf(ctx, segmentID, segmentType, scope, topic, collectionID, ts) + if rf, ok := ret.Get(0).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID, int64) []*model.Segment); ok { + r0 = rf(ctx, segmentID, segmentType, scope, collectionID, ts) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]*model.Segment) } } - if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, *string, *string, types.UniqueID, int64) error); ok { - r1 = rf(ctx, segmentID, segmentType, scope, topic, collectionID, ts) + if rf, ok := ret.Get(1).(func(context.Context, types.UniqueID, *string, *string, types.UniqueID, int64) error); ok { + r1 = rf(ctx, segmentID, segmentType, scope, collectionID, ts) } else { r1 = ret.Error(1) } diff --git a/go/coordinator/internal/model/collection.go b/go/pkg/model/collection.go similarity index 70% rename from go/coordinator/internal/model/collection.go rename to go/pkg/model/collection.go index 6e242b7fc67..ec35569daab 100644 --- a/go/coordinator/internal/model/collection.go +++ b/go/pkg/model/collection.go @@ -1,25 +1,24 @@ package model import ( - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/types" ) type Collection struct { ID types.UniqueID Name string - Topic string Dimension *int32 Metadata *CollectionMetadata[CollectionMetadataValueType] - Created bool TenantID string DatabaseName string Ts types.Timestamp + LogPosition int64 + Version int32 } type CreateCollection struct { ID types.UniqueID Name string - Topic string Dimension *int32 Metadata *CollectionMetadata[CollectionMetadataValueType] GetOrCreate bool @@ -38,7 +37,6 @@ type DeleteCollection struct { type UpdateCollection struct { ID types.UniqueID Name *string - Topic *string Dimension *int32 Metadata *CollectionMetadata[CollectionMetadataValueType] ResetMetadata bool @@ -47,15 +45,26 @@ type UpdateCollection struct { Ts types.Timestamp } -func FilterCollection(collection *Collection, collectionID types.UniqueID, collectionName *string, collectionTopic *string) bool { +type FlushCollectionCompaction struct { + ID types.UniqueID + TenantID string + LogPosition int64 + CurrentCollectionVersion int32 + FlushSegmentCompactions []*FlushSegmentCompaction +} + +type FlushCollectionInfo struct { + ID string + CollectionVersion int32 + TenantLastCompactionTime int64 +} + +func FilterCollection(collection *Collection, collectionID types.UniqueID, collectionName *string) bool { if collectionID != types.NilUniqueID() && collectionID != collection.ID { return false } if collectionName != nil && *collectionName != collection.Name { return false } - if collectionTopic != nil && *collectionTopic != collection.Topic { - return false - } return true } diff --git a/go/coordinator/internal/model/collection_metadata.go b/go/pkg/model/collection_metadata.go similarity index 100% rename from go/coordinator/internal/model/collection_metadata.go rename to go/pkg/model/collection_metadata.go diff --git a/go/coordinator/internal/model/database.go b/go/pkg/model/database.go similarity index 83% rename from go/coordinator/internal/model/database.go rename to go/pkg/model/database.go index ad23e3f14c6..debbbecb21d 100644 --- a/go/coordinator/internal/model/database.go +++ b/go/pkg/model/database.go @@ -1,6 +1,6 @@ package model -import "github.com/chroma/chroma-coordinator/internal/types" +import "github.com/chroma-core/chroma/go/pkg/types" type Database struct { ID string diff --git a/go/coordinator/internal/model/notification.go b/go/pkg/model/notification.go similarity index 100% rename from go/coordinator/internal/model/notification.go rename to go/pkg/model/notification.go diff --git a/go/coordinator/internal/model/segment.go b/go/pkg/model/segment.go similarity index 86% rename from go/coordinator/internal/model/segment.go rename to go/pkg/model/segment.go index 8fa93f10cca..5e30c96df1c 100644 --- a/go/coordinator/internal/model/segment.go +++ b/go/pkg/model/segment.go @@ -1,24 +1,23 @@ package model import ( - "github.com/chroma/chroma-coordinator/internal/types" + "github.com/chroma-core/chroma/go/pkg/types" ) type Segment struct { ID types.UniqueID Type string Scope string - Topic *string CollectionID types.UniqueID Metadata *SegmentMetadata[SegmentMetadataValueType] Ts types.Timestamp + FilePaths map[string][]string } type CreateSegment struct { ID types.UniqueID Type string Scope string - Topic *string CollectionID types.UniqueID Metadata *SegmentMetadata[SegmentMetadataValueType] Ts types.Timestamp @@ -26,7 +25,6 @@ type CreateSegment struct { type UpdateSegment struct { ID types.UniqueID - Topic *string ResetTopic bool Collection *string ResetCollection bool @@ -39,10 +37,14 @@ type GetSegments struct { ID types.UniqueID Type *string Scope *string - Topic *string CollectionID types.UniqueID } +type FlushSegmentCompaction struct { + ID types.UniqueID + FilePaths map[string][]string +} + func FilterSegments(segment *Segment, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) bool { if segmentID != types.NilUniqueID() && segment.ID != segmentID { return false @@ -55,10 +57,6 @@ func FilterSegments(segment *Segment, segmentID types.UniqueID, segmentType *str return false } - if topic != nil && *segment.Topic != *topic { - return false - } - if collectionID != types.NilUniqueID() && segment.CollectionID != collectionID { return false } diff --git a/go/coordinator/internal/model/segment_metadata.go b/go/pkg/model/segment_metadata.go similarity index 100% rename from go/coordinator/internal/model/segment_metadata.go rename to go/pkg/model/segment_metadata.go diff --git a/go/coordinator/internal/model/tenant.go b/go/pkg/model/tenant.go similarity index 58% rename from go/coordinator/internal/model/tenant.go rename to go/pkg/model/tenant.go index 191d781d00a..ac7131eb81c 100644 --- a/go/coordinator/internal/model/tenant.go +++ b/go/pkg/model/tenant.go @@ -1,6 +1,6 @@ package model -import "github.com/chroma/chroma-coordinator/internal/types" +import "github.com/chroma-core/chroma/go/pkg/types" type Tenant struct { Name string @@ -15,3 +15,8 @@ type GetTenant struct { Name string Ts types.Timestamp } + +type TenantLastCompactionTime struct { + ID string + Ts types.Timestamp +} diff --git a/go/coordinator/internal/notification/database_notification_store.go b/go/pkg/notification/database_notification_store.go similarity index 95% rename from go/coordinator/internal/notification/database_notification_store.go rename to go/pkg/notification/database_notification_store.go index 93411ff3973..9f0b5f4bc38 100644 --- a/go/coordinator/internal/notification/database_notification_store.go +++ b/go/pkg/notification/database_notification_store.go @@ -4,8 +4,8 @@ import ( "context" "sort" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/model" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/model" ) type DatabaseNotificationStore struct { diff --git a/go/coordinator/internal/notification/database_notification_store_test.go b/go/pkg/notification/database_notification_store_test.go similarity index 97% rename from go/coordinator/internal/notification/database_notification_store_test.go rename to go/pkg/notification/database_notification_store_test.go index d2e9fa91f2c..6d91fd2cb3a 100644 --- a/go/coordinator/internal/notification/database_notification_store_test.go +++ b/go/pkg/notification/database_notification_store_test.go @@ -5,9 +5,9 @@ import ( "reflect" "testing" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel/mocks" - "github.com/chroma/chroma-coordinator/internal/model" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel/mocks" + "github.com/chroma-core/chroma/go/pkg/model" "github.com/stretchr/testify/mock" ) diff --git a/go/coordinator/internal/notification/memory_notification_store.go b/go/pkg/notification/memory_notification_store.go similarity index 97% rename from go/coordinator/internal/notification/memory_notification_store.go rename to go/pkg/notification/memory_notification_store.go index e6168d9a3ca..80beba1d537 100644 --- a/go/coordinator/internal/notification/memory_notification_store.go +++ b/go/pkg/notification/memory_notification_store.go @@ -4,7 +4,7 @@ import ( "context" "sort" - "github.com/chroma/chroma-coordinator/internal/model" + "github.com/chroma-core/chroma/go/pkg/model" ) type MemoryNotificationStore struct { diff --git a/go/coordinator/internal/notification/memory_notification_store_test.go b/go/pkg/notification/memory_notification_store_test.go similarity index 98% rename from go/coordinator/internal/notification/memory_notification_store_test.go rename to go/pkg/notification/memory_notification_store_test.go index 17734898f3d..24b3c08c044 100644 --- a/go/coordinator/internal/notification/memory_notification_store_test.go +++ b/go/pkg/notification/memory_notification_store_test.go @@ -5,7 +5,7 @@ import ( "reflect" "testing" - "github.com/chroma/chroma-coordinator/internal/model" + "github.com/chroma-core/chroma/go/pkg/model" ) func TestMemoryNotificationStore_GetAllPendingNotifications(t *testing.T) { diff --git a/go/coordinator/internal/notification/notification_processor.go b/go/pkg/notification/notification_processor.go similarity index 97% rename from go/coordinator/internal/notification/notification_processor.go rename to go/pkg/notification/notification_processor.go index e9113dc000d..32fd51f16af 100644 --- a/go/coordinator/internal/notification/notification_processor.go +++ b/go/pkg/notification/notification_processor.go @@ -4,8 +4,8 @@ import ( "context" "sync/atomic" - "github.com/chroma/chroma-coordinator/internal/common" - "github.com/chroma/chroma-coordinator/internal/model" + "github.com/chroma-core/chroma/go/pkg/common" + "github.com/chroma-core/chroma/go/pkg/model" "github.com/pingcap/log" "go.uber.org/zap" ) diff --git a/go/coordinator/internal/notification/notification_processor_test.go b/go/pkg/notification/notification_processor_test.go similarity index 92% rename from go/coordinator/internal/notification/notification_processor_test.go rename to go/pkg/notification/notification_processor_test.go index 23c85d27eb8..7efab440049 100644 --- a/go/coordinator/internal/notification/notification_processor_test.go +++ b/go/pkg/notification/notification_processor_test.go @@ -4,11 +4,11 @@ import ( "context" "testing" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dao" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbcore" - "github.com/chroma/chroma-coordinator/internal/metastore/db/dbmodel" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dao" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbcore" + "github.com/chroma-core/chroma/go/pkg/metastore/db/dbmodel" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" "google.golang.org/protobuf/proto" "gorm.io/driver/sqlite" "gorm.io/gorm" diff --git a/go/coordinator/internal/notification/notification_store.go b/go/pkg/notification/notification_store.go similarity index 88% rename from go/coordinator/internal/notification/notification_store.go rename to go/pkg/notification/notification_store.go index 6e0434ffa5b..60ef61278eb 100644 --- a/go/coordinator/internal/notification/notification_store.go +++ b/go/pkg/notification/notification_store.go @@ -3,7 +3,7 @@ package notification import ( "context" - "github.com/chroma/chroma-coordinator/internal/model" + "github.com/chroma-core/chroma/go/pkg/model" ) type NotificationStore interface { diff --git a/go/coordinator/internal/notification/notifier.go b/go/pkg/notification/notifier.go similarity index 95% rename from go/coordinator/internal/notification/notifier.go rename to go/pkg/notification/notifier.go index ce19bb62c5e..62ace26d8a9 100644 --- a/go/coordinator/internal/notification/notifier.go +++ b/go/pkg/notification/notifier.go @@ -4,8 +4,8 @@ import ( "context" "github.com/apache/pulsar-client-go/pulsar" - "github.com/chroma/chroma-coordinator/internal/model" - "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb" + "github.com/chroma-core/chroma/go/pkg/model" + "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" "github.com/pingcap/log" "go.uber.org/zap" "google.golang.org/protobuf/proto" diff --git a/go/coordinator/internal/proto/coordinatorpb/chroma.pb.go b/go/pkg/proto/coordinatorpb/chroma.pb.go similarity index 67% rename from go/coordinator/internal/proto/coordinatorpb/chroma.pb.go rename to go/pkg/proto/coordinatorpb/chroma.pb.go index 3cec5eefe06..53dcc20607e 100644 --- a/go/coordinator/internal/proto/coordinatorpb/chroma.pb.go +++ b/go/pkg/proto/coordinatorpb/chroma.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.31.0 +// protoc-gen-go v1.33.0 // protoc v4.23.4 // source: chromadb/proto/chroma.proto @@ -20,6 +20,7 @@ const ( _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) ) +// Types here should mirror chromadb/types.py type Operation int32 const ( @@ -219,16 +220,18 @@ func (x *Status) GetCode() int32 { return 0 } -type ChromaResponse struct { +type Vector struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` + Dimension int32 `protobuf:"varint,1,opt,name=dimension,proto3" json:"dimension,omitempty"` + Vector []byte `protobuf:"bytes,2,opt,name=vector,proto3" json:"vector,omitempty"` + Encoding ScalarEncoding `protobuf:"varint,3,opt,name=encoding,proto3,enum=chroma.ScalarEncoding" json:"encoding,omitempty"` } -func (x *ChromaResponse) Reset() { - *x = ChromaResponse{} +func (x *Vector) Reset() { + *x = Vector{} if protoimpl.UnsafeEnabled { mi := &file_chromadb_proto_chroma_proto_msgTypes[1] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -236,13 +239,13 @@ func (x *ChromaResponse) Reset() { } } -func (x *ChromaResponse) String() string { +func (x *Vector) String() string { return protoimpl.X.MessageStringOf(x) } -func (*ChromaResponse) ProtoMessage() {} +func (*Vector) ProtoMessage() {} -func (x *ChromaResponse) ProtoReflect() protoreflect.Message { +func (x *Vector) ProtoReflect() protoreflect.Message { mi := &file_chromadb_proto_chroma_proto_msgTypes[1] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -254,30 +257,42 @@ func (x *ChromaResponse) ProtoReflect() protoreflect.Message { return mi.MessageOf(x) } -// Deprecated: Use ChromaResponse.ProtoReflect.Descriptor instead. -func (*ChromaResponse) Descriptor() ([]byte, []int) { +// Deprecated: Use Vector.ProtoReflect.Descriptor instead. +func (*Vector) Descriptor() ([]byte, []int) { return file_chromadb_proto_chroma_proto_rawDescGZIP(), []int{1} } -func (x *ChromaResponse) GetStatus() *Status { +func (x *Vector) GetDimension() int32 { + if x != nil { + return x.Dimension + } + return 0 +} + +func (x *Vector) GetVector() []byte { if x != nil { - return x.Status + return x.Vector } return nil } -type Vector struct { +func (x *Vector) GetEncoding() ScalarEncoding { + if x != nil { + return x.Encoding + } + return ScalarEncoding_FLOAT32 +} + +type FilePaths struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - Dimension int32 `protobuf:"varint,1,opt,name=dimension,proto3" json:"dimension,omitempty"` - Vector []byte `protobuf:"bytes,2,opt,name=vector,proto3" json:"vector,omitempty"` - Encoding ScalarEncoding `protobuf:"varint,3,opt,name=encoding,proto3,enum=chroma.ScalarEncoding" json:"encoding,omitempty"` + Paths []string `protobuf:"bytes,1,rep,name=paths,proto3" json:"paths,omitempty"` } -func (x *Vector) Reset() { - *x = Vector{} +func (x *FilePaths) Reset() { + *x = FilePaths{} if protoimpl.UnsafeEnabled { mi := &file_chromadb_proto_chroma_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -285,13 +300,13 @@ func (x *Vector) Reset() { } } -func (x *Vector) String() string { +func (x *FilePaths) String() string { return protoimpl.X.MessageStringOf(x) } -func (*Vector) ProtoMessage() {} +func (*FilePaths) ProtoMessage() {} -func (x *Vector) ProtoReflect() protoreflect.Message { +func (x *FilePaths) ProtoReflect() protoreflect.Message { mi := &file_chromadb_proto_chroma_proto_msgTypes[2] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -303,45 +318,29 @@ func (x *Vector) ProtoReflect() protoreflect.Message { return mi.MessageOf(x) } -// Deprecated: Use Vector.ProtoReflect.Descriptor instead. -func (*Vector) Descriptor() ([]byte, []int) { +// Deprecated: Use FilePaths.ProtoReflect.Descriptor instead. +func (*FilePaths) Descriptor() ([]byte, []int) { return file_chromadb_proto_chroma_proto_rawDescGZIP(), []int{2} } -func (x *Vector) GetDimension() int32 { +func (x *FilePaths) GetPaths() []string { if x != nil { - return x.Dimension - } - return 0 -} - -func (x *Vector) GetVector() []byte { - if x != nil { - return x.Vector + return x.Paths } return nil } -func (x *Vector) GetEncoding() ScalarEncoding { - if x != nil { - return x.Encoding - } - return ScalarEncoding_FLOAT32 -} - type Segment struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Type string `protobuf:"bytes,2,opt,name=type,proto3" json:"type,omitempty"` - Scope SegmentScope `protobuf:"varint,3,opt,name=scope,proto3,enum=chroma.SegmentScope" json:"scope,omitempty"` - Topic *string `protobuf:"bytes,4,opt,name=topic,proto3,oneof" json:"topic,omitempty"` // TODO should channel <> segment binding exist here? - // If a segment has a collection, it implies that this segment implements the full - // collection and can be used to service queries (for it's given scope.) - Collection *string `protobuf:"bytes,5,opt,name=collection,proto3,oneof" json:"collection,omitempty"` - Metadata *UpdateMetadata `protobuf:"bytes,6,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Type string `protobuf:"bytes,2,opt,name=type,proto3" json:"type,omitempty"` + Scope SegmentScope `protobuf:"varint,3,opt,name=scope,proto3,enum=chroma.SegmentScope" json:"scope,omitempty"` + Collection *string `protobuf:"bytes,5,opt,name=collection,proto3,oneof" json:"collection,omitempty"` + Metadata *UpdateMetadata `protobuf:"bytes,6,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` + FilePaths map[string]*FilePaths `protobuf:"bytes,7,rep,name=file_paths,json=filePaths,proto3" json:"file_paths,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` } func (x *Segment) Reset() { @@ -397,13 +396,6 @@ func (x *Segment) GetScope() SegmentScope { return SegmentScope_VECTOR } -func (x *Segment) GetTopic() string { - if x != nil && x.Topic != nil { - return *x.Topic - } - return "" -} - func (x *Segment) GetCollection() string { if x != nil && x.Collection != nil { return *x.Collection @@ -418,18 +410,26 @@ func (x *Segment) GetMetadata() *UpdateMetadata { return nil } +func (x *Segment) GetFilePaths() map[string]*FilePaths { + if x != nil { + return x.FilePaths + } + return nil +} + type Collection struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` - Topic string `protobuf:"bytes,3,opt,name=topic,proto3" json:"topic,omitempty"` - Metadata *UpdateMetadata `protobuf:"bytes,4,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` - Dimension *int32 `protobuf:"varint,5,opt,name=dimension,proto3,oneof" json:"dimension,omitempty"` - Tenant string `protobuf:"bytes,6,opt,name=tenant,proto3" json:"tenant,omitempty"` - Database string `protobuf:"bytes,7,opt,name=database,proto3" json:"database,omitempty"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` + Metadata *UpdateMetadata `protobuf:"bytes,4,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` + Dimension *int32 `protobuf:"varint,5,opt,name=dimension,proto3,oneof" json:"dimension,omitempty"` + Tenant string `protobuf:"bytes,6,opt,name=tenant,proto3" json:"tenant,omitempty"` + Database string `protobuf:"bytes,7,opt,name=database,proto3" json:"database,omitempty"` + LogPosition int64 `protobuf:"varint,8,opt,name=log_position,json=logPosition,proto3" json:"log_position,omitempty"` + Version int32 `protobuf:"varint,9,opt,name=version,proto3" json:"version,omitempty"` } func (x *Collection) Reset() { @@ -478,13 +478,6 @@ func (x *Collection) GetName() string { return "" } -func (x *Collection) GetTopic() string { - if x != nil { - return x.Topic - } - return "" -} - func (x *Collection) GetMetadata() *UpdateMetadata { if x != nil { return x.Metadata @@ -513,6 +506,20 @@ func (x *Collection) GetDatabase() string { return "" } +func (x *Collection) GetLogPosition() int64 { + if x != nil { + return x.LogPosition + } + return 0 +} + +func (x *Collection) GetVersion() int32 { + if x != nil { + return x.Version + } + return 0 +} + type Database struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache @@ -765,20 +772,20 @@ func (x *UpdateMetadata) GetMetadata() map[string]*UpdateMetadataValue { return nil } -type SubmitEmbeddingRecord struct { +// Represents an operation the user submits +type OperationRecord struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Vector *Vector `protobuf:"bytes,2,opt,name=vector,proto3,oneof" json:"vector,omitempty"` - Metadata *UpdateMetadata `protobuf:"bytes,3,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` - Operation Operation `protobuf:"varint,4,opt,name=operation,proto3,enum=chroma.Operation" json:"operation,omitempty"` - CollectionId string `protobuf:"bytes,5,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Vector *Vector `protobuf:"bytes,2,opt,name=vector,proto3,oneof" json:"vector,omitempty"` + Metadata *UpdateMetadata `protobuf:"bytes,3,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` + Operation Operation `protobuf:"varint,4,opt,name=operation,proto3,enum=chroma.Operation" json:"operation,omitempty"` } -func (x *SubmitEmbeddingRecord) Reset() { - *x = SubmitEmbeddingRecord{} +func (x *OperationRecord) Reset() { + *x = OperationRecord{} if protoimpl.UnsafeEnabled { mi := &file_chromadb_proto_chroma_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -786,13 +793,13 @@ func (x *SubmitEmbeddingRecord) Reset() { } } -func (x *SubmitEmbeddingRecord) String() string { +func (x *OperationRecord) String() string { return protoimpl.X.MessageStringOf(x) } -func (*SubmitEmbeddingRecord) ProtoMessage() {} +func (*OperationRecord) ProtoMessage() {} -func (x *SubmitEmbeddingRecord) ProtoReflect() protoreflect.Message { +func (x *OperationRecord) ProtoReflect() protoreflect.Message { mi := &file_chromadb_proto_chroma_proto_msgTypes[9] if protoimpl.UnsafeEnabled && x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -804,53 +811,45 @@ func (x *SubmitEmbeddingRecord) ProtoReflect() protoreflect.Message { return mi.MessageOf(x) } -// Deprecated: Use SubmitEmbeddingRecord.ProtoReflect.Descriptor instead. -func (*SubmitEmbeddingRecord) Descriptor() ([]byte, []int) { +// Deprecated: Use OperationRecord.ProtoReflect.Descriptor instead. +func (*OperationRecord) Descriptor() ([]byte, []int) { return file_chromadb_proto_chroma_proto_rawDescGZIP(), []int{9} } -func (x *SubmitEmbeddingRecord) GetId() string { +func (x *OperationRecord) GetId() string { if x != nil { return x.Id } return "" } -func (x *SubmitEmbeddingRecord) GetVector() *Vector { +func (x *OperationRecord) GetVector() *Vector { if x != nil { return x.Vector } return nil } -func (x *SubmitEmbeddingRecord) GetMetadata() *UpdateMetadata { +func (x *OperationRecord) GetMetadata() *UpdateMetadata { if x != nil { return x.Metadata } return nil } -func (x *SubmitEmbeddingRecord) GetOperation() Operation { +func (x *OperationRecord) GetOperation() Operation { if x != nil { return x.Operation } return Operation_ADD } -func (x *SubmitEmbeddingRecord) GetCollectionId() string { - if x != nil { - return x.CollectionId - } - return "" -} - type VectorEmbeddingRecord struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - SeqId []byte `protobuf:"bytes,2,opt,name=seq_id,json=seqId,proto3" json:"seq_id,omitempty"` Vector *Vector `protobuf:"bytes,3,opt,name=vector,proto3" json:"vector,omitempty"` // TODO: we need to rethink source of truth for vector dimensionality and encoding } @@ -893,13 +892,6 @@ func (x *VectorEmbeddingRecord) GetId() string { return "" } -func (x *VectorEmbeddingRecord) GetSeqId() []byte { - if x != nil { - return x.SeqId - } - return nil -} - func (x *VectorEmbeddingRecord) GetVector() *Vector { if x != nil { return x.Vector @@ -913,8 +905,7 @@ type VectorQueryResult struct { unknownFields protoimpl.UnknownFields Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - SeqId []byte `protobuf:"bytes,2,opt,name=seq_id,json=seqId,proto3" json:"seq_id,omitempty"` - Distance float64 `protobuf:"fixed64,3,opt,name=distance,proto3" json:"distance,omitempty"` + Distance float32 `protobuf:"fixed32,3,opt,name=distance,proto3" json:"distance,omitempty"` Vector *Vector `protobuf:"bytes,4,opt,name=vector,proto3,oneof" json:"vector,omitempty"` } @@ -957,14 +948,7 @@ func (x *VectorQueryResult) GetId() string { return "" } -func (x *VectorQueryResult) GetSeqId() []byte { - if x != nil { - return x.SeqId - } - return nil -} - -func (x *VectorQueryResult) GetDistance() float64 { +func (x *VectorQueryResult) GetDistance() float32 { if x != nil { return x.Distance } @@ -1261,161 +1245,163 @@ var file_chromadb_proto_chroma_proto_rawDesc = []byte{ 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x22, 0x34, 0x0a, 0x06, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x16, 0x0a, 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x12, 0x12, 0x0a, 0x04, 0x63, 0x6f, 0x64, 0x65, 0x18, - 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x63, 0x6f, 0x64, 0x65, 0x22, 0x38, 0x0a, 0x0e, 0x43, - 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, - 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, - 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, - 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x72, 0x0a, 0x06, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x12, - 0x1c, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, - 0x28, 0x05, 0x52, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x16, 0x0a, - 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x06, 0x76, - 0x65, 0x63, 0x74, 0x6f, 0x72, 0x12, 0x32, 0x0a, 0x08, 0x65, 0x6e, 0x63, 0x6f, 0x64, 0x69, 0x6e, - 0x67, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, - 0x2e, 0x53, 0x63, 0x61, 0x6c, 0x61, 0x72, 0x45, 0x6e, 0x63, 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x52, - 0x08, 0x65, 0x6e, 0x63, 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x22, 0xf8, 0x01, 0x0a, 0x07, 0x53, 0x65, - 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x2a, 0x0a, 0x05, 0x73, 0x63, 0x6f, - 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x14, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x53, 0x63, 0x6f, 0x70, 0x65, 0x52, 0x05, - 0x73, 0x63, 0x6f, 0x70, 0x65, 0x12, 0x19, 0x0a, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x18, 0x04, - 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, 0x52, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x88, 0x01, 0x01, - 0x12, 0x23, 0x0a, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x05, - 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, - 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x37, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, - 0x61, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, - 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, - 0x02, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x88, 0x01, 0x01, 0x42, 0x08, - 0x0a, 0x06, 0x5f, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x42, 0x0d, 0x0a, 0x0b, 0x5f, 0x63, 0x6f, 0x6c, - 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x42, 0x0b, 0x0a, 0x09, 0x5f, 0x6d, 0x65, 0x74, 0x61, - 0x64, 0x61, 0x74, 0x61, 0x22, 0xf1, 0x01, 0x0a, 0x0a, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, - 0x69, 0x6f, 0x6e, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x14, 0x0a, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, - 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x12, 0x37, 0x0a, - 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, - 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, - 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x00, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, - 0x61, 0x74, 0x61, 0x88, 0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, - 0x69, 0x6f, 0x6e, 0x18, 0x05, 0x20, 0x01, 0x28, 0x05, 0x48, 0x01, 0x52, 0x09, 0x64, 0x69, 0x6d, - 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, - 0x61, 0x6e, 0x74, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, - 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x07, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x42, 0x0b, 0x0a, - 0x09, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x64, - 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x22, 0x46, 0x0a, 0x08, 0x44, 0x61, 0x74, 0x61, - 0x62, 0x61, 0x73, 0x65, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, - 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, - 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, - 0x6e, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, - 0x22, 0x1c, 0x0a, 0x06, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, - 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x85, - 0x01, 0x0a, 0x13, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, - 0x61, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x12, 0x23, 0x0a, 0x0c, 0x73, 0x74, 0x72, 0x69, 0x6e, 0x67, - 0x5f, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, 0x52, 0x0b, - 0x73, 0x74, 0x72, 0x69, 0x6e, 0x67, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x12, 0x1d, 0x0a, 0x09, 0x69, - 0x6e, 0x74, 0x5f, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x48, 0x00, - 0x52, 0x08, 0x69, 0x6e, 0x74, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x12, 0x21, 0x0a, 0x0b, 0x66, 0x6c, - 0x6f, 0x61, 0x74, 0x5f, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x01, 0x48, - 0x00, 0x52, 0x0a, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x42, 0x07, 0x0a, - 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x22, 0xac, 0x01, 0x0a, 0x0e, 0x55, 0x70, 0x64, 0x61, 0x74, - 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x12, 0x40, 0x0a, 0x08, 0x6d, 0x65, 0x74, - 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x24, 0x2e, 0x63, 0x68, - 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, - 0x61, 0x74, 0x61, 0x2e, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, - 0x79, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x1a, 0x58, 0x0a, 0x0d, 0x4d, - 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, - 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x31, - 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1b, 0x2e, + 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x63, 0x6f, 0x64, 0x65, 0x22, 0x72, 0x0a, 0x06, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x12, 0x1c, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, + 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, + 0x69, 0x6f, 0x6e, 0x12, 0x16, 0x0a, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x02, 0x20, + 0x01, 0x28, 0x0c, 0x52, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x12, 0x32, 0x0a, 0x08, 0x65, + 0x6e, 0x63, 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x16, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x63, 0x61, 0x6c, 0x61, 0x72, 0x45, 0x6e, 0x63, + 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x08, 0x65, 0x6e, 0x63, 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x22, + 0x21, 0x0a, 0x09, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x12, 0x14, 0x0a, 0x05, + 0x70, 0x61, 0x74, 0x68, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, 0x05, 0x70, 0x61, 0x74, + 0x68, 0x73, 0x22, 0xe3, 0x02, 0x0a, 0x07, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x0e, + 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, + 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, + 0x70, 0x65, 0x12, 0x2a, 0x0a, 0x05, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, + 0x0e, 0x32, 0x14, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, + 0x6e, 0x74, 0x53, 0x63, 0x6f, 0x70, 0x65, 0x52, 0x05, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x12, 0x23, + 0x0a, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x05, 0x20, 0x01, + 0x28, 0x09, 0x48, 0x00, 0x52, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x88, 0x01, 0x01, 0x12, 0x37, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, + 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, + 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x01, 0x52, + 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x88, 0x01, 0x01, 0x12, 0x3d, 0x0a, 0x0a, + 0x66, 0x69, 0x6c, 0x65, 0x5f, 0x70, 0x61, 0x74, 0x68, 0x73, 0x18, 0x07, 0x20, 0x03, 0x28, 0x0b, + 0x32, 0x1e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, + 0x74, 0x2e, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, + 0x52, 0x09, 0x66, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x1a, 0x4f, 0x0a, 0x0e, 0x46, + 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, + 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, + 0x27, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, + 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, + 0x73, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x42, 0x0d, 0x0a, 0x0b, + 0x5f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x42, 0x0b, 0x0a, 0x09, 0x5f, + 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x22, 0x98, 0x02, 0x0a, 0x0a, 0x43, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, + 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x37, 0x0a, 0x08, 0x6d, + 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, - 0x61, 0x64, 0x61, 0x74, 0x61, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, - 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0xfb, 0x01, 0x0a, 0x15, 0x53, 0x75, 0x62, 0x6d, 0x69, 0x74, - 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x12, + 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x00, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, + 0x61, 0x88, 0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, + 0x6e, 0x18, 0x05, 0x20, 0x01, 0x28, 0x05, 0x48, 0x01, 0x52, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, + 0x73, 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, + 0x74, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, + 0x1a, 0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x07, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x12, 0x21, 0x0a, 0x0c, 0x6c, + 0x6f, 0x67, 0x5f, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x08, 0x20, 0x01, 0x28, + 0x03, 0x52, 0x0b, 0x6c, 0x6f, 0x67, 0x50, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x18, + 0x0a, 0x07, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x09, 0x20, 0x01, 0x28, 0x05, 0x52, + 0x07, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x42, 0x0b, 0x0a, 0x09, 0x5f, 0x6d, 0x65, 0x74, + 0x61, 0x64, 0x61, 0x74, 0x61, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, + 0x69, 0x6f, 0x6e, 0x22, 0x46, 0x0a, 0x08, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, - 0x2b, 0x0a, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, - 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x48, - 0x00, 0x52, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x88, 0x01, 0x01, 0x12, 0x37, 0x0a, 0x08, - 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, - 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, - 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x01, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, - 0x74, 0x61, 0x88, 0x01, 0x01, 0x12, 0x2f, 0x0a, 0x09, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, - 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x11, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x2e, 0x4f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x09, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, - 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, - 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x42, 0x09, 0x0a, 0x07, 0x5f, - 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x42, 0x0b, 0x0a, 0x09, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, - 0x61, 0x74, 0x61, 0x22, 0x66, 0x0a, 0x15, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x45, 0x6d, 0x62, - 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x12, 0x0e, 0x0a, 0x02, - 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x15, 0x0a, 0x06, - 0x73, 0x65, 0x71, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x05, 0x73, 0x65, - 0x71, 0x49, 0x64, 0x12, 0x26, 0x0a, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x03, 0x20, - 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, 0x63, - 0x74, 0x6f, 0x72, 0x52, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x22, 0x8e, 0x01, 0x0a, 0x11, - 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, - 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, - 0x64, 0x12, 0x15, 0x0a, 0x06, 0x73, 0x65, 0x71, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, - 0x0c, 0x52, 0x05, 0x73, 0x65, 0x71, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x69, 0x73, 0x74, - 0x61, 0x6e, 0x63, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x01, 0x52, 0x08, 0x64, 0x69, 0x73, 0x74, - 0x61, 0x6e, 0x63, 0x65, 0x12, 0x2b, 0x0a, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x04, + 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, + 0x61, 0x6d, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x03, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x22, 0x1c, 0x0a, 0x06, 0x54, + 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x85, 0x01, 0x0a, 0x13, 0x55, 0x70, + 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x56, 0x61, 0x6c, 0x75, + 0x65, 0x12, 0x23, 0x0a, 0x0c, 0x73, 0x74, 0x72, 0x69, 0x6e, 0x67, 0x5f, 0x76, 0x61, 0x6c, 0x75, + 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, 0x52, 0x0b, 0x73, 0x74, 0x72, 0x69, 0x6e, + 0x67, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x12, 0x1d, 0x0a, 0x09, 0x69, 0x6e, 0x74, 0x5f, 0x76, 0x61, + 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x48, 0x00, 0x52, 0x08, 0x69, 0x6e, 0x74, + 0x56, 0x61, 0x6c, 0x75, 0x65, 0x12, 0x21, 0x0a, 0x0b, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x5f, 0x76, + 0x61, 0x6c, 0x75, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x01, 0x48, 0x00, 0x52, 0x0a, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x56, 0x61, 0x6c, 0x75, 0x65, 0x42, 0x07, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, + 0x65, 0x22, 0xac, 0x01, 0x0a, 0x0e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, + 0x64, 0x61, 0x74, 0x61, 0x12, 0x40, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, + 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x24, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x2e, 0x4d, + 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x08, 0x6d, 0x65, + 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x1a, 0x58, 0x0a, 0x0d, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, + 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x31, 0x0a, 0x05, 0x76, 0x61, 0x6c, + 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, + 0x56, 0x61, 0x6c, 0x75, 0x65, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, + 0x22, 0xd0, 0x01, 0x0a, 0x0f, 0x4f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, + 0x52, 0x02, 0x69, 0x64, 0x12, 0x2b, 0x0a, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x48, 0x00, 0x52, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x88, 0x01, - 0x01, 0x42, 0x09, 0x0a, 0x07, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x22, 0x49, 0x0a, 0x12, - 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, - 0x74, 0x73, 0x12, 0x33, 0x0a, 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x18, 0x01, 0x20, - 0x03, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, 0x63, - 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x07, - 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x22, 0x44, 0x0a, 0x11, 0x47, 0x65, 0x74, 0x56, 0x65, - 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x10, 0x0a, 0x03, - 0x69, 0x64, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, 0x03, 0x69, 0x64, 0x73, 0x12, 0x1d, - 0x0a, 0x0a, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, - 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x22, 0x4d, 0x0a, - 0x12, 0x47, 0x65, 0x74, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x12, 0x37, 0x0a, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x18, 0x01, - 0x20, 0x03, 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, - 0x63, 0x74, 0x6f, 0x72, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x65, 0x63, - 0x6f, 0x72, 0x64, 0x52, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x22, 0xbc, 0x01, 0x0a, - 0x13, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x12, 0x28, 0x0a, 0x07, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x18, - 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, - 0x65, 0x63, 0x74, 0x6f, 0x72, 0x52, 0x07, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x12, 0x0c, - 0x0a, 0x01, 0x6b, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x01, 0x6b, 0x12, 0x1f, 0x0a, 0x0b, - 0x61, 0x6c, 0x6c, 0x6f, 0x77, 0x65, 0x64, 0x5f, 0x69, 0x64, 0x73, 0x18, 0x03, 0x20, 0x03, 0x28, - 0x09, 0x52, 0x0a, 0x61, 0x6c, 0x6c, 0x6f, 0x77, 0x65, 0x64, 0x49, 0x64, 0x73, 0x12, 0x2d, 0x0a, - 0x12, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x5f, 0x65, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, - 0x6e, 0x67, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x11, 0x69, 0x6e, 0x63, 0x6c, 0x75, - 0x64, 0x65, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x73, 0x12, 0x1d, 0x0a, 0x0a, - 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, - 0x52, 0x09, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x22, 0x4c, 0x0a, 0x14, 0x51, - 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x12, 0x34, 0x0a, 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x18, 0x01, - 0x20, 0x03, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, - 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, - 0x52, 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x2a, 0x38, 0x0a, 0x09, 0x4f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x07, 0x0a, 0x03, 0x41, 0x44, 0x44, 0x10, 0x00, 0x12, - 0x0a, 0x0a, 0x06, 0x55, 0x50, 0x44, 0x41, 0x54, 0x45, 0x10, 0x01, 0x12, 0x0a, 0x0a, 0x06, 0x55, - 0x50, 0x53, 0x45, 0x52, 0x54, 0x10, 0x02, 0x12, 0x0a, 0x0a, 0x06, 0x44, 0x45, 0x4c, 0x45, 0x54, - 0x45, 0x10, 0x03, 0x2a, 0x28, 0x0a, 0x0e, 0x53, 0x63, 0x61, 0x6c, 0x61, 0x72, 0x45, 0x6e, 0x63, - 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x12, 0x0b, 0x0a, 0x07, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x33, 0x32, - 0x10, 0x00, 0x12, 0x09, 0x0a, 0x05, 0x49, 0x4e, 0x54, 0x33, 0x32, 0x10, 0x01, 0x2a, 0x28, 0x0a, - 0x0c, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x53, 0x63, 0x6f, 0x70, 0x65, 0x12, 0x0a, 0x0a, - 0x06, 0x56, 0x45, 0x43, 0x54, 0x4f, 0x52, 0x10, 0x00, 0x12, 0x0c, 0x0a, 0x08, 0x4d, 0x45, 0x54, - 0x41, 0x44, 0x41, 0x54, 0x41, 0x10, 0x01, 0x32, 0xa2, 0x01, 0x0a, 0x0c, 0x56, 0x65, 0x63, 0x74, - 0x6f, 0x72, 0x52, 0x65, 0x61, 0x64, 0x65, 0x72, 0x12, 0x45, 0x0a, 0x0a, 0x47, 0x65, 0x74, 0x56, - 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x12, 0x19, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, - 0x47, 0x65, 0x74, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x56, 0x65, - 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, - 0x4b, 0x0a, 0x0c, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x12, - 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, - 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x63, - 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, - 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x43, 0x5a, 0x41, - 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, - 0x61, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2d, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, - 0x61, 0x74, 0x6f, 0x72, 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x70, 0x72, - 0x6f, 0x74, 0x6f, 0x2f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x70, - 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x01, 0x12, 0x37, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x03, 0x20, + 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, + 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x01, 0x52, 0x08, 0x6d, + 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x88, 0x01, 0x01, 0x12, 0x2f, 0x0a, 0x09, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x11, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x4f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, + 0x52, 0x09, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x42, 0x09, 0x0a, 0x07, 0x5f, + 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x42, 0x0b, 0x0a, 0x09, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, + 0x61, 0x74, 0x61, 0x22, 0x4f, 0x0a, 0x15, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x45, 0x6d, 0x62, + 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x12, 0x0e, 0x0a, 0x02, + 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x26, 0x0a, 0x06, + 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, + 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x52, 0x06, 0x76, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x22, 0x77, 0x0a, 0x11, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, + 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x69, 0x73, + 0x74, 0x61, 0x6e, 0x63, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x02, 0x52, 0x08, 0x64, 0x69, 0x73, + 0x74, 0x61, 0x6e, 0x63, 0x65, 0x12, 0x2b, 0x0a, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x18, + 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x48, 0x00, 0x52, 0x06, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x88, + 0x01, 0x01, 0x42, 0x09, 0x0a, 0x07, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x22, 0x49, 0x0a, + 0x12, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, + 0x6c, 0x74, 0x73, 0x12, 0x33, 0x0a, 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x18, 0x01, + 0x20, 0x03, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, + 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x22, 0x44, 0x0a, 0x11, 0x47, 0x65, 0x74, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x10, 0x0a, + 0x03, 0x69, 0x64, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, 0x03, 0x69, 0x64, 0x73, 0x12, + 0x1d, 0x0a, 0x0a, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x22, 0x4d, + 0x0a, 0x12, 0x47, 0x65, 0x74, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, + 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x37, 0x0a, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x18, + 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x52, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x22, 0xbc, 0x01, + 0x0a, 0x13, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x28, 0x0a, 0x07, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, + 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x52, 0x07, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x12, + 0x0c, 0x0a, 0x01, 0x6b, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x01, 0x6b, 0x12, 0x1f, 0x0a, + 0x0b, 0x61, 0x6c, 0x6c, 0x6f, 0x77, 0x65, 0x64, 0x5f, 0x69, 0x64, 0x73, 0x18, 0x03, 0x20, 0x03, + 0x28, 0x09, 0x52, 0x0a, 0x61, 0x6c, 0x6c, 0x6f, 0x77, 0x65, 0x64, 0x49, 0x64, 0x73, 0x12, 0x2d, + 0x0a, 0x12, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x5f, 0x65, 0x6d, 0x62, 0x65, 0x64, 0x64, + 0x69, 0x6e, 0x67, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x11, 0x69, 0x6e, 0x63, 0x6c, + 0x75, 0x64, 0x65, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x73, 0x12, 0x1d, 0x0a, + 0x0a, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x05, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x09, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x22, 0x4c, 0x0a, 0x14, + 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, + 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x34, 0x0a, 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x18, + 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, + 0x73, 0x52, 0x07, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x73, 0x2a, 0x38, 0x0a, 0x09, 0x4f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x07, 0x0a, 0x03, 0x41, 0x44, 0x44, 0x10, 0x00, + 0x12, 0x0a, 0x0a, 0x06, 0x55, 0x50, 0x44, 0x41, 0x54, 0x45, 0x10, 0x01, 0x12, 0x0a, 0x0a, 0x06, + 0x55, 0x50, 0x53, 0x45, 0x52, 0x54, 0x10, 0x02, 0x12, 0x0a, 0x0a, 0x06, 0x44, 0x45, 0x4c, 0x45, + 0x54, 0x45, 0x10, 0x03, 0x2a, 0x28, 0x0a, 0x0e, 0x53, 0x63, 0x61, 0x6c, 0x61, 0x72, 0x45, 0x6e, + 0x63, 0x6f, 0x64, 0x69, 0x6e, 0x67, 0x12, 0x0b, 0x0a, 0x07, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x33, + 0x32, 0x10, 0x00, 0x12, 0x09, 0x0a, 0x05, 0x49, 0x4e, 0x54, 0x33, 0x32, 0x10, 0x01, 0x2a, 0x28, + 0x0a, 0x0c, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x53, 0x63, 0x6f, 0x70, 0x65, 0x12, 0x0a, + 0x0a, 0x06, 0x56, 0x45, 0x43, 0x54, 0x4f, 0x52, 0x10, 0x00, 0x12, 0x0c, 0x0a, 0x08, 0x4d, 0x45, + 0x54, 0x41, 0x44, 0x41, 0x54, 0x41, 0x10, 0x01, 0x32, 0xa2, 0x01, 0x0a, 0x0c, 0x56, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x52, 0x65, 0x61, 0x64, 0x65, 0x72, 0x12, 0x45, 0x0a, 0x0a, 0x47, 0x65, 0x74, + 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x12, 0x19, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x47, 0x65, 0x74, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, + 0x12, 0x4b, 0x0a, 0x0c, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, + 0x12, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, + 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x56, 0x65, 0x63, 0x74, + 0x6f, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x3a, 0x5a, + 0x38, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2d, 0x63, 0x6f, 0x72, 0x65, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2f, 0x67, + 0x6f, 0x2f, 0x70, 0x6b, 0x67, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x33, } var ( @@ -1431,21 +1417,21 @@ func file_chromadb_proto_chroma_proto_rawDescGZIP() []byte { } var file_chromadb_proto_chroma_proto_enumTypes = make([]protoimpl.EnumInfo, 3) -var file_chromadb_proto_chroma_proto_msgTypes = make([]protoimpl.MessageInfo, 18) +var file_chromadb_proto_chroma_proto_msgTypes = make([]protoimpl.MessageInfo, 19) var file_chromadb_proto_chroma_proto_goTypes = []interface{}{ (Operation)(0), // 0: chroma.Operation (ScalarEncoding)(0), // 1: chroma.ScalarEncoding (SegmentScope)(0), // 2: chroma.SegmentScope (*Status)(nil), // 3: chroma.Status - (*ChromaResponse)(nil), // 4: chroma.ChromaResponse - (*Vector)(nil), // 5: chroma.Vector + (*Vector)(nil), // 4: chroma.Vector + (*FilePaths)(nil), // 5: chroma.FilePaths (*Segment)(nil), // 6: chroma.Segment (*Collection)(nil), // 7: chroma.Collection (*Database)(nil), // 8: chroma.Database (*Tenant)(nil), // 9: chroma.Tenant (*UpdateMetadataValue)(nil), // 10: chroma.UpdateMetadataValue (*UpdateMetadata)(nil), // 11: chroma.UpdateMetadata - (*SubmitEmbeddingRecord)(nil), // 12: chroma.SubmitEmbeddingRecord + (*OperationRecord)(nil), // 12: chroma.OperationRecord (*VectorEmbeddingRecord)(nil), // 13: chroma.VectorEmbeddingRecord (*VectorQueryResult)(nil), // 14: chroma.VectorQueryResult (*VectorQueryResults)(nil), // 15: chroma.VectorQueryResults @@ -1453,34 +1439,36 @@ var file_chromadb_proto_chroma_proto_goTypes = []interface{}{ (*GetVectorsResponse)(nil), // 17: chroma.GetVectorsResponse (*QueryVectorsRequest)(nil), // 18: chroma.QueryVectorsRequest (*QueryVectorsResponse)(nil), // 19: chroma.QueryVectorsResponse - nil, // 20: chroma.UpdateMetadata.MetadataEntry + nil, // 20: chroma.Segment.FilePathsEntry + nil, // 21: chroma.UpdateMetadata.MetadataEntry } var file_chromadb_proto_chroma_proto_depIdxs = []int32{ - 3, // 0: chroma.ChromaResponse.status:type_name -> chroma.Status - 1, // 1: chroma.Vector.encoding:type_name -> chroma.ScalarEncoding - 2, // 2: chroma.Segment.scope:type_name -> chroma.SegmentScope - 11, // 3: chroma.Segment.metadata:type_name -> chroma.UpdateMetadata + 1, // 0: chroma.Vector.encoding:type_name -> chroma.ScalarEncoding + 2, // 1: chroma.Segment.scope:type_name -> chroma.SegmentScope + 11, // 2: chroma.Segment.metadata:type_name -> chroma.UpdateMetadata + 20, // 3: chroma.Segment.file_paths:type_name -> chroma.Segment.FilePathsEntry 11, // 4: chroma.Collection.metadata:type_name -> chroma.UpdateMetadata - 20, // 5: chroma.UpdateMetadata.metadata:type_name -> chroma.UpdateMetadata.MetadataEntry - 5, // 6: chroma.SubmitEmbeddingRecord.vector:type_name -> chroma.Vector - 11, // 7: chroma.SubmitEmbeddingRecord.metadata:type_name -> chroma.UpdateMetadata - 0, // 8: chroma.SubmitEmbeddingRecord.operation:type_name -> chroma.Operation - 5, // 9: chroma.VectorEmbeddingRecord.vector:type_name -> chroma.Vector - 5, // 10: chroma.VectorQueryResult.vector:type_name -> chroma.Vector + 21, // 5: chroma.UpdateMetadata.metadata:type_name -> chroma.UpdateMetadata.MetadataEntry + 4, // 6: chroma.OperationRecord.vector:type_name -> chroma.Vector + 11, // 7: chroma.OperationRecord.metadata:type_name -> chroma.UpdateMetadata + 0, // 8: chroma.OperationRecord.operation:type_name -> chroma.Operation + 4, // 9: chroma.VectorEmbeddingRecord.vector:type_name -> chroma.Vector + 4, // 10: chroma.VectorQueryResult.vector:type_name -> chroma.Vector 14, // 11: chroma.VectorQueryResults.results:type_name -> chroma.VectorQueryResult 13, // 12: chroma.GetVectorsResponse.records:type_name -> chroma.VectorEmbeddingRecord - 5, // 13: chroma.QueryVectorsRequest.vectors:type_name -> chroma.Vector + 4, // 13: chroma.QueryVectorsRequest.vectors:type_name -> chroma.Vector 15, // 14: chroma.QueryVectorsResponse.results:type_name -> chroma.VectorQueryResults - 10, // 15: chroma.UpdateMetadata.MetadataEntry.value:type_name -> chroma.UpdateMetadataValue - 16, // 16: chroma.VectorReader.GetVectors:input_type -> chroma.GetVectorsRequest - 18, // 17: chroma.VectorReader.QueryVectors:input_type -> chroma.QueryVectorsRequest - 17, // 18: chroma.VectorReader.GetVectors:output_type -> chroma.GetVectorsResponse - 19, // 19: chroma.VectorReader.QueryVectors:output_type -> chroma.QueryVectorsResponse - 18, // [18:20] is the sub-list for method output_type - 16, // [16:18] is the sub-list for method input_type - 16, // [16:16] is the sub-list for extension type_name - 16, // [16:16] is the sub-list for extension extendee - 0, // [0:16] is the sub-list for field type_name + 5, // 15: chroma.Segment.FilePathsEntry.value:type_name -> chroma.FilePaths + 10, // 16: chroma.UpdateMetadata.MetadataEntry.value:type_name -> chroma.UpdateMetadataValue + 16, // 17: chroma.VectorReader.GetVectors:input_type -> chroma.GetVectorsRequest + 18, // 18: chroma.VectorReader.QueryVectors:input_type -> chroma.QueryVectorsRequest + 17, // 19: chroma.VectorReader.GetVectors:output_type -> chroma.GetVectorsResponse + 19, // 20: chroma.VectorReader.QueryVectors:output_type -> chroma.QueryVectorsResponse + 19, // [19:21] is the sub-list for method output_type + 17, // [17:19] is the sub-list for method input_type + 17, // [17:17] is the sub-list for extension type_name + 17, // [17:17] is the sub-list for extension extendee + 0, // [0:17] is the sub-list for field type_name } func init() { file_chromadb_proto_chroma_proto_init() } @@ -1502,7 +1490,7 @@ func file_chromadb_proto_chroma_proto_init() { } } file_chromadb_proto_chroma_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*ChromaResponse); i { + switch v := v.(*Vector); i { case 0: return &v.state case 1: @@ -1514,7 +1502,7 @@ func file_chromadb_proto_chroma_proto_init() { } } file_chromadb_proto_chroma_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*Vector); i { + switch v := v.(*FilePaths); i { case 0: return &v.state case 1: @@ -1598,7 +1586,7 @@ func file_chromadb_proto_chroma_proto_init() { } } file_chromadb_proto_chroma_proto_msgTypes[9].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*SubmitEmbeddingRecord); i { + switch v := v.(*OperationRecord); i { case 0: return &v.state case 1: @@ -1709,7 +1697,7 @@ func file_chromadb_proto_chroma_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: file_chromadb_proto_chroma_proto_rawDesc, NumEnums: 3, - NumMessages: 18, + NumMessages: 19, NumExtensions: 0, NumServices: 1, }, diff --git a/go/coordinator/internal/proto/coordinatorpb/chroma_grpc.pb.go b/go/pkg/proto/coordinatorpb/chroma_grpc.pb.go similarity index 100% rename from go/coordinator/internal/proto/coordinatorpb/chroma_grpc.pb.go rename to go/pkg/proto/coordinatorpb/chroma_grpc.pb.go diff --git a/go/pkg/proto/coordinatorpb/coordinator.pb.go b/go/pkg/proto/coordinatorpb/coordinator.pb.go new file mode 100644 index 00000000000..2b75f87cfa8 --- /dev/null +++ b/go/pkg/proto/coordinatorpb/coordinator.pb.go @@ -0,0 +1,2898 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.33.0 +// protoc v4.23.4 +// source: chromadb/proto/coordinator.proto + +package coordinatorpb + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + emptypb "google.golang.org/protobuf/types/known/emptypb" + reflect "reflect" + sync "sync" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type CreateDatabaseRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` + Tenant string `protobuf:"bytes,3,opt,name=tenant,proto3" json:"tenant,omitempty"` +} + +func (x *CreateDatabaseRequest) Reset() { + *x = CreateDatabaseRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateDatabaseRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateDatabaseRequest) ProtoMessage() {} + +func (x *CreateDatabaseRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[0] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateDatabaseRequest.ProtoReflect.Descriptor instead. +func (*CreateDatabaseRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{0} +} + +func (x *CreateDatabaseRequest) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *CreateDatabaseRequest) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *CreateDatabaseRequest) GetTenant() string { + if x != nil { + return x.Tenant + } + return "" +} + +type CreateDatabaseResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *CreateDatabaseResponse) Reset() { + *x = CreateDatabaseResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateDatabaseResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateDatabaseResponse) ProtoMessage() {} + +func (x *CreateDatabaseResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[1] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateDatabaseResponse.ProtoReflect.Descriptor instead. +func (*CreateDatabaseResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{1} +} + +func (x *CreateDatabaseResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type GetDatabaseRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` + Tenant string `protobuf:"bytes,2,opt,name=tenant,proto3" json:"tenant,omitempty"` +} + +func (x *GetDatabaseRequest) Reset() { + *x = GetDatabaseRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetDatabaseRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetDatabaseRequest) ProtoMessage() {} + +func (x *GetDatabaseRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[2] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetDatabaseRequest.ProtoReflect.Descriptor instead. +func (*GetDatabaseRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{2} +} + +func (x *GetDatabaseRequest) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *GetDatabaseRequest) GetTenant() string { + if x != nil { + return x.Tenant + } + return "" +} + +type GetDatabaseResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Database *Database `protobuf:"bytes,1,opt,name=database,proto3" json:"database,omitempty"` + Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *GetDatabaseResponse) Reset() { + *x = GetDatabaseResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetDatabaseResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetDatabaseResponse) ProtoMessage() {} + +func (x *GetDatabaseResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[3] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetDatabaseResponse.ProtoReflect.Descriptor instead. +func (*GetDatabaseResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{3} +} + +func (x *GetDatabaseResponse) GetDatabase() *Database { + if x != nil { + return x.Database + } + return nil +} + +func (x *GetDatabaseResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type CreateTenantRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` // Names are globally unique +} + +func (x *CreateTenantRequest) Reset() { + *x = CreateTenantRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateTenantRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateTenantRequest) ProtoMessage() {} + +func (x *CreateTenantRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[4] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateTenantRequest.ProtoReflect.Descriptor instead. +func (*CreateTenantRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{4} +} + +func (x *CreateTenantRequest) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +type CreateTenantResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *CreateTenantResponse) Reset() { + *x = CreateTenantResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateTenantResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateTenantResponse) ProtoMessage() {} + +func (x *CreateTenantResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[5] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateTenantResponse.ProtoReflect.Descriptor instead. +func (*CreateTenantResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{5} +} + +func (x *CreateTenantResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type GetTenantRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` +} + +func (x *GetTenantRequest) Reset() { + *x = GetTenantRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetTenantRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetTenantRequest) ProtoMessage() {} + +func (x *GetTenantRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[6] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetTenantRequest.ProtoReflect.Descriptor instead. +func (*GetTenantRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{6} +} + +func (x *GetTenantRequest) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +type GetTenantResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Tenant *Tenant `protobuf:"bytes,1,opt,name=tenant,proto3" json:"tenant,omitempty"` + Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *GetTenantResponse) Reset() { + *x = GetTenantResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetTenantResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetTenantResponse) ProtoMessage() {} + +func (x *GetTenantResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[7] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetTenantResponse.ProtoReflect.Descriptor instead. +func (*GetTenantResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{7} +} + +func (x *GetTenantResponse) GetTenant() *Tenant { + if x != nil { + return x.Tenant + } + return nil +} + +func (x *GetTenantResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type CreateSegmentRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Segment *Segment `protobuf:"bytes,1,opt,name=segment,proto3" json:"segment,omitempty"` +} + +func (x *CreateSegmentRequest) Reset() { + *x = CreateSegmentRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateSegmentRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateSegmentRequest) ProtoMessage() {} + +func (x *CreateSegmentRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[8] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateSegmentRequest.ProtoReflect.Descriptor instead. +func (*CreateSegmentRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{8} +} + +func (x *CreateSegmentRequest) GetSegment() *Segment { + if x != nil { + return x.Segment + } + return nil +} + +type CreateSegmentResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *CreateSegmentResponse) Reset() { + *x = CreateSegmentResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateSegmentResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateSegmentResponse) ProtoMessage() {} + +func (x *CreateSegmentResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[9] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateSegmentResponse.ProtoReflect.Descriptor instead. +func (*CreateSegmentResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{9} +} + +func (x *CreateSegmentResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type DeleteSegmentRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` +} + +func (x *DeleteSegmentRequest) Reset() { + *x = DeleteSegmentRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *DeleteSegmentRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteSegmentRequest) ProtoMessage() {} + +func (x *DeleteSegmentRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[10] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteSegmentRequest.ProtoReflect.Descriptor instead. +func (*DeleteSegmentRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{10} +} + +func (x *DeleteSegmentRequest) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +type DeleteSegmentResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *DeleteSegmentResponse) Reset() { + *x = DeleteSegmentResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *DeleteSegmentResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteSegmentResponse) ProtoMessage() {} + +func (x *DeleteSegmentResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[11] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteSegmentResponse.ProtoReflect.Descriptor instead. +func (*DeleteSegmentResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{11} +} + +func (x *DeleteSegmentResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type GetSegmentsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id *string `protobuf:"bytes,1,opt,name=id,proto3,oneof" json:"id,omitempty"` + Type *string `protobuf:"bytes,2,opt,name=type,proto3,oneof" json:"type,omitempty"` + Scope *SegmentScope `protobuf:"varint,3,opt,name=scope,proto3,enum=chroma.SegmentScope,oneof" json:"scope,omitempty"` + Collection *string `protobuf:"bytes,5,opt,name=collection,proto3,oneof" json:"collection,omitempty"` // Collection ID +} + +func (x *GetSegmentsRequest) Reset() { + *x = GetSegmentsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetSegmentsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetSegmentsRequest) ProtoMessage() {} + +func (x *GetSegmentsRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[12] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetSegmentsRequest.ProtoReflect.Descriptor instead. +func (*GetSegmentsRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{12} +} + +func (x *GetSegmentsRequest) GetId() string { + if x != nil && x.Id != nil { + return *x.Id + } + return "" +} + +func (x *GetSegmentsRequest) GetType() string { + if x != nil && x.Type != nil { + return *x.Type + } + return "" +} + +func (x *GetSegmentsRequest) GetScope() SegmentScope { + if x != nil && x.Scope != nil { + return *x.Scope + } + return SegmentScope_VECTOR +} + +func (x *GetSegmentsRequest) GetCollection() string { + if x != nil && x.Collection != nil { + return *x.Collection + } + return "" +} + +type GetSegmentsResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Segments []*Segment `protobuf:"bytes,1,rep,name=segments,proto3" json:"segments,omitempty"` + Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *GetSegmentsResponse) Reset() { + *x = GetSegmentsResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetSegmentsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetSegmentsResponse) ProtoMessage() {} + +func (x *GetSegmentsResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[13] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetSegmentsResponse.ProtoReflect.Descriptor instead. +func (*GetSegmentsResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{13} +} + +func (x *GetSegmentsResponse) GetSegments() []*Segment { + if x != nil { + return x.Segments + } + return nil +} + +func (x *GetSegmentsResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type UpdateSegmentRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + // Types that are assignable to CollectionUpdate: + // + // *UpdateSegmentRequest_Collection + // *UpdateSegmentRequest_ResetCollection + CollectionUpdate isUpdateSegmentRequest_CollectionUpdate `protobuf_oneof:"collection_update"` + // Types that are assignable to MetadataUpdate: + // + // *UpdateSegmentRequest_Metadata + // *UpdateSegmentRequest_ResetMetadata + MetadataUpdate isUpdateSegmentRequest_MetadataUpdate `protobuf_oneof:"metadata_update"` +} + +func (x *UpdateSegmentRequest) Reset() { + *x = UpdateSegmentRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *UpdateSegmentRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpdateSegmentRequest) ProtoMessage() {} + +func (x *UpdateSegmentRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[14] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpdateSegmentRequest.ProtoReflect.Descriptor instead. +func (*UpdateSegmentRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{14} +} + +func (x *UpdateSegmentRequest) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (m *UpdateSegmentRequest) GetCollectionUpdate() isUpdateSegmentRequest_CollectionUpdate { + if m != nil { + return m.CollectionUpdate + } + return nil +} + +func (x *UpdateSegmentRequest) GetCollection() string { + if x, ok := x.GetCollectionUpdate().(*UpdateSegmentRequest_Collection); ok { + return x.Collection + } + return "" +} + +func (x *UpdateSegmentRequest) GetResetCollection() bool { + if x, ok := x.GetCollectionUpdate().(*UpdateSegmentRequest_ResetCollection); ok { + return x.ResetCollection + } + return false +} + +func (m *UpdateSegmentRequest) GetMetadataUpdate() isUpdateSegmentRequest_MetadataUpdate { + if m != nil { + return m.MetadataUpdate + } + return nil +} + +func (x *UpdateSegmentRequest) GetMetadata() *UpdateMetadata { + if x, ok := x.GetMetadataUpdate().(*UpdateSegmentRequest_Metadata); ok { + return x.Metadata + } + return nil +} + +func (x *UpdateSegmentRequest) GetResetMetadata() bool { + if x, ok := x.GetMetadataUpdate().(*UpdateSegmentRequest_ResetMetadata); ok { + return x.ResetMetadata + } + return false +} + +type isUpdateSegmentRequest_CollectionUpdate interface { + isUpdateSegmentRequest_CollectionUpdate() +} + +type UpdateSegmentRequest_Collection struct { + Collection string `protobuf:"bytes,4,opt,name=collection,proto3,oneof"` +} + +type UpdateSegmentRequest_ResetCollection struct { + ResetCollection bool `protobuf:"varint,5,opt,name=reset_collection,json=resetCollection,proto3,oneof"` +} + +func (*UpdateSegmentRequest_Collection) isUpdateSegmentRequest_CollectionUpdate() {} + +func (*UpdateSegmentRequest_ResetCollection) isUpdateSegmentRequest_CollectionUpdate() {} + +type isUpdateSegmentRequest_MetadataUpdate interface { + isUpdateSegmentRequest_MetadataUpdate() +} + +type UpdateSegmentRequest_Metadata struct { + Metadata *UpdateMetadata `protobuf:"bytes,6,opt,name=metadata,proto3,oneof"` +} + +type UpdateSegmentRequest_ResetMetadata struct { + ResetMetadata bool `protobuf:"varint,7,opt,name=reset_metadata,json=resetMetadata,proto3,oneof"` +} + +func (*UpdateSegmentRequest_Metadata) isUpdateSegmentRequest_MetadataUpdate() {} + +func (*UpdateSegmentRequest_ResetMetadata) isUpdateSegmentRequest_MetadataUpdate() {} + +type UpdateSegmentResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *UpdateSegmentResponse) Reset() { + *x = UpdateSegmentResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *UpdateSegmentResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpdateSegmentResponse) ProtoMessage() {} + +func (x *UpdateSegmentResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[15] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpdateSegmentResponse.ProtoReflect.Descriptor instead. +func (*UpdateSegmentResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{15} +} + +func (x *UpdateSegmentResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type CreateCollectionRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` + Metadata *UpdateMetadata `protobuf:"bytes,3,opt,name=metadata,proto3,oneof" json:"metadata,omitempty"` + Dimension *int32 `protobuf:"varint,4,opt,name=dimension,proto3,oneof" json:"dimension,omitempty"` + GetOrCreate *bool `protobuf:"varint,5,opt,name=get_or_create,json=getOrCreate,proto3,oneof" json:"get_or_create,omitempty"` + Tenant string `protobuf:"bytes,6,opt,name=tenant,proto3" json:"tenant,omitempty"` + Database string `protobuf:"bytes,7,opt,name=database,proto3" json:"database,omitempty"` +} + +func (x *CreateCollectionRequest) Reset() { + *x = CreateCollectionRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateCollectionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateCollectionRequest) ProtoMessage() {} + +func (x *CreateCollectionRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[16] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateCollectionRequest.ProtoReflect.Descriptor instead. +func (*CreateCollectionRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{16} +} + +func (x *CreateCollectionRequest) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *CreateCollectionRequest) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *CreateCollectionRequest) GetMetadata() *UpdateMetadata { + if x != nil { + return x.Metadata + } + return nil +} + +func (x *CreateCollectionRequest) GetDimension() int32 { + if x != nil && x.Dimension != nil { + return *x.Dimension + } + return 0 +} + +func (x *CreateCollectionRequest) GetGetOrCreate() bool { + if x != nil && x.GetOrCreate != nil { + return *x.GetOrCreate + } + return false +} + +func (x *CreateCollectionRequest) GetTenant() string { + if x != nil { + return x.Tenant + } + return "" +} + +func (x *CreateCollectionRequest) GetDatabase() string { + if x != nil { + return x.Database + } + return "" +} + +type CreateCollectionResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Collection *Collection `protobuf:"bytes,1,opt,name=collection,proto3" json:"collection,omitempty"` + Created bool `protobuf:"varint,2,opt,name=created,proto3" json:"created,omitempty"` + Status *Status `protobuf:"bytes,3,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *CreateCollectionResponse) Reset() { + *x = CreateCollectionResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CreateCollectionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CreateCollectionResponse) ProtoMessage() {} + +func (x *CreateCollectionResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[17] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CreateCollectionResponse.ProtoReflect.Descriptor instead. +func (*CreateCollectionResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{17} +} + +func (x *CreateCollectionResponse) GetCollection() *Collection { + if x != nil { + return x.Collection + } + return nil +} + +func (x *CreateCollectionResponse) GetCreated() bool { + if x != nil { + return x.Created + } + return false +} + +func (x *CreateCollectionResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type DeleteCollectionRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Tenant string `protobuf:"bytes,2,opt,name=tenant,proto3" json:"tenant,omitempty"` + Database string `protobuf:"bytes,3,opt,name=database,proto3" json:"database,omitempty"` +} + +func (x *DeleteCollectionRequest) Reset() { + *x = DeleteCollectionRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[18] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *DeleteCollectionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteCollectionRequest) ProtoMessage() {} + +func (x *DeleteCollectionRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[18] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteCollectionRequest.ProtoReflect.Descriptor instead. +func (*DeleteCollectionRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{18} +} + +func (x *DeleteCollectionRequest) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *DeleteCollectionRequest) GetTenant() string { + if x != nil { + return x.Tenant + } + return "" +} + +func (x *DeleteCollectionRequest) GetDatabase() string { + if x != nil { + return x.Database + } + return "" +} + +type DeleteCollectionResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *DeleteCollectionResponse) Reset() { + *x = DeleteCollectionResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[19] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *DeleteCollectionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteCollectionResponse) ProtoMessage() {} + +func (x *DeleteCollectionResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[19] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteCollectionResponse.ProtoReflect.Descriptor instead. +func (*DeleteCollectionResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{19} +} + +func (x *DeleteCollectionResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type GetCollectionsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id *string `protobuf:"bytes,1,opt,name=id,proto3,oneof" json:"id,omitempty"` + Name *string `protobuf:"bytes,2,opt,name=name,proto3,oneof" json:"name,omitempty"` + Tenant string `protobuf:"bytes,4,opt,name=tenant,proto3" json:"tenant,omitempty"` + Database string `protobuf:"bytes,5,opt,name=database,proto3" json:"database,omitempty"` +} + +func (x *GetCollectionsRequest) Reset() { + *x = GetCollectionsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[20] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetCollectionsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetCollectionsRequest) ProtoMessage() {} + +func (x *GetCollectionsRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[20] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetCollectionsRequest.ProtoReflect.Descriptor instead. +func (*GetCollectionsRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{20} +} + +func (x *GetCollectionsRequest) GetId() string { + if x != nil && x.Id != nil { + return *x.Id + } + return "" +} + +func (x *GetCollectionsRequest) GetName() string { + if x != nil && x.Name != nil { + return *x.Name + } + return "" +} + +func (x *GetCollectionsRequest) GetTenant() string { + if x != nil { + return x.Tenant + } + return "" +} + +func (x *GetCollectionsRequest) GetDatabase() string { + if x != nil { + return x.Database + } + return "" +} + +type GetCollectionsResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Collections []*Collection `protobuf:"bytes,1,rep,name=collections,proto3" json:"collections,omitempty"` + Status *Status `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *GetCollectionsResponse) Reset() { + *x = GetCollectionsResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[21] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetCollectionsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetCollectionsResponse) ProtoMessage() {} + +func (x *GetCollectionsResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[21] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetCollectionsResponse.ProtoReflect.Descriptor instead. +func (*GetCollectionsResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{21} +} + +func (x *GetCollectionsResponse) GetCollections() []*Collection { + if x != nil { + return x.Collections + } + return nil +} + +func (x *GetCollectionsResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type UpdateCollectionRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Name *string `protobuf:"bytes,3,opt,name=name,proto3,oneof" json:"name,omitempty"` + Dimension *int32 `protobuf:"varint,4,opt,name=dimension,proto3,oneof" json:"dimension,omitempty"` + // Types that are assignable to MetadataUpdate: + // + // *UpdateCollectionRequest_Metadata + // *UpdateCollectionRequest_ResetMetadata + MetadataUpdate isUpdateCollectionRequest_MetadataUpdate `protobuf_oneof:"metadata_update"` +} + +func (x *UpdateCollectionRequest) Reset() { + *x = UpdateCollectionRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[22] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *UpdateCollectionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpdateCollectionRequest) ProtoMessage() {} + +func (x *UpdateCollectionRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[22] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpdateCollectionRequest.ProtoReflect.Descriptor instead. +func (*UpdateCollectionRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{22} +} + +func (x *UpdateCollectionRequest) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *UpdateCollectionRequest) GetName() string { + if x != nil && x.Name != nil { + return *x.Name + } + return "" +} + +func (x *UpdateCollectionRequest) GetDimension() int32 { + if x != nil && x.Dimension != nil { + return *x.Dimension + } + return 0 +} + +func (m *UpdateCollectionRequest) GetMetadataUpdate() isUpdateCollectionRequest_MetadataUpdate { + if m != nil { + return m.MetadataUpdate + } + return nil +} + +func (x *UpdateCollectionRequest) GetMetadata() *UpdateMetadata { + if x, ok := x.GetMetadataUpdate().(*UpdateCollectionRequest_Metadata); ok { + return x.Metadata + } + return nil +} + +func (x *UpdateCollectionRequest) GetResetMetadata() bool { + if x, ok := x.GetMetadataUpdate().(*UpdateCollectionRequest_ResetMetadata); ok { + return x.ResetMetadata + } + return false +} + +type isUpdateCollectionRequest_MetadataUpdate interface { + isUpdateCollectionRequest_MetadataUpdate() +} + +type UpdateCollectionRequest_Metadata struct { + Metadata *UpdateMetadata `protobuf:"bytes,5,opt,name=metadata,proto3,oneof"` +} + +type UpdateCollectionRequest_ResetMetadata struct { + ResetMetadata bool `protobuf:"varint,6,opt,name=reset_metadata,json=resetMetadata,proto3,oneof"` +} + +func (*UpdateCollectionRequest_Metadata) isUpdateCollectionRequest_MetadataUpdate() {} + +func (*UpdateCollectionRequest_ResetMetadata) isUpdateCollectionRequest_MetadataUpdate() {} + +type UpdateCollectionResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *UpdateCollectionResponse) Reset() { + *x = UpdateCollectionResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[23] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *UpdateCollectionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpdateCollectionResponse) ProtoMessage() {} + +func (x *UpdateCollectionResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[23] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpdateCollectionResponse.ProtoReflect.Descriptor instead. +func (*UpdateCollectionResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{23} +} + +func (x *UpdateCollectionResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type Notification struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Id int64 `protobuf:"varint,1,opt,name=id,proto3" json:"id,omitempty"` + CollectionId string `protobuf:"bytes,2,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + Type string `protobuf:"bytes,3,opt,name=type,proto3" json:"type,omitempty"` + Status string `protobuf:"bytes,4,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *Notification) Reset() { + *x = Notification{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[24] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *Notification) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Notification) ProtoMessage() {} + +func (x *Notification) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[24] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Notification.ProtoReflect.Descriptor instead. +func (*Notification) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{24} +} + +func (x *Notification) GetId() int64 { + if x != nil { + return x.Id + } + return 0 +} + +func (x *Notification) GetCollectionId() string { + if x != nil { + return x.CollectionId + } + return "" +} + +func (x *Notification) GetType() string { + if x != nil { + return x.Type + } + return "" +} + +func (x *Notification) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +type ResetStateResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Status *Status `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` +} + +func (x *ResetStateResponse) Reset() { + *x = ResetStateResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[25] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *ResetStateResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ResetStateResponse) ProtoMessage() {} + +func (x *ResetStateResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[25] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ResetStateResponse.ProtoReflect.Descriptor instead. +func (*ResetStateResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{25} +} + +func (x *ResetStateResponse) GetStatus() *Status { + if x != nil { + return x.Status + } + return nil +} + +type GetLastCompactionTimeForTenantRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + TenantId []string `protobuf:"bytes,1,rep,name=tenant_id,json=tenantId,proto3" json:"tenant_id,omitempty"` +} + +func (x *GetLastCompactionTimeForTenantRequest) Reset() { + *x = GetLastCompactionTimeForTenantRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[26] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetLastCompactionTimeForTenantRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetLastCompactionTimeForTenantRequest) ProtoMessage() {} + +func (x *GetLastCompactionTimeForTenantRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[26] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetLastCompactionTimeForTenantRequest.ProtoReflect.Descriptor instead. +func (*GetLastCompactionTimeForTenantRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{26} +} + +func (x *GetLastCompactionTimeForTenantRequest) GetTenantId() []string { + if x != nil { + return x.TenantId + } + return nil +} + +type TenantLastCompactionTime struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + TenantId string `protobuf:"bytes,1,opt,name=tenant_id,json=tenantId,proto3" json:"tenant_id,omitempty"` + LastCompactionTime int64 `protobuf:"varint,2,opt,name=last_compaction_time,json=lastCompactionTime,proto3" json:"last_compaction_time,omitempty"` +} + +func (x *TenantLastCompactionTime) Reset() { + *x = TenantLastCompactionTime{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[27] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *TenantLastCompactionTime) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TenantLastCompactionTime) ProtoMessage() {} + +func (x *TenantLastCompactionTime) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[27] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TenantLastCompactionTime.ProtoReflect.Descriptor instead. +func (*TenantLastCompactionTime) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{27} +} + +func (x *TenantLastCompactionTime) GetTenantId() string { + if x != nil { + return x.TenantId + } + return "" +} + +func (x *TenantLastCompactionTime) GetLastCompactionTime() int64 { + if x != nil { + return x.LastCompactionTime + } + return 0 +} + +type GetLastCompactionTimeForTenantResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + TenantLastCompactionTime []*TenantLastCompactionTime `protobuf:"bytes,1,rep,name=tenant_last_compaction_time,json=tenantLastCompactionTime,proto3" json:"tenant_last_compaction_time,omitempty"` +} + +func (x *GetLastCompactionTimeForTenantResponse) Reset() { + *x = GetLastCompactionTimeForTenantResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[28] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetLastCompactionTimeForTenantResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetLastCompactionTimeForTenantResponse) ProtoMessage() {} + +func (x *GetLastCompactionTimeForTenantResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[28] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetLastCompactionTimeForTenantResponse.ProtoReflect.Descriptor instead. +func (*GetLastCompactionTimeForTenantResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{28} +} + +func (x *GetLastCompactionTimeForTenantResponse) GetTenantLastCompactionTime() []*TenantLastCompactionTime { + if x != nil { + return x.TenantLastCompactionTime + } + return nil +} + +type SetLastCompactionTimeForTenantRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + TenantLastCompactionTime *TenantLastCompactionTime `protobuf:"bytes,1,opt,name=tenant_last_compaction_time,json=tenantLastCompactionTime,proto3" json:"tenant_last_compaction_time,omitempty"` +} + +func (x *SetLastCompactionTimeForTenantRequest) Reset() { + *x = SetLastCompactionTimeForTenantRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[29] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *SetLastCompactionTimeForTenantRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SetLastCompactionTimeForTenantRequest) ProtoMessage() {} + +func (x *SetLastCompactionTimeForTenantRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[29] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SetLastCompactionTimeForTenantRequest.ProtoReflect.Descriptor instead. +func (*SetLastCompactionTimeForTenantRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{29} +} + +func (x *SetLastCompactionTimeForTenantRequest) GetTenantLastCompactionTime() *TenantLastCompactionTime { + if x != nil { + return x.TenantLastCompactionTime + } + return nil +} + +type FlushSegmentCompactionInfo struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + SegmentId string `protobuf:"bytes,1,opt,name=segment_id,json=segmentId,proto3" json:"segment_id,omitempty"` + FilePaths map[string]*FilePaths `protobuf:"bytes,2,rep,name=file_paths,json=filePaths,proto3" json:"file_paths,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` +} + +func (x *FlushSegmentCompactionInfo) Reset() { + *x = FlushSegmentCompactionInfo{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[30] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *FlushSegmentCompactionInfo) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FlushSegmentCompactionInfo) ProtoMessage() {} + +func (x *FlushSegmentCompactionInfo) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[30] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FlushSegmentCompactionInfo.ProtoReflect.Descriptor instead. +func (*FlushSegmentCompactionInfo) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{30} +} + +func (x *FlushSegmentCompactionInfo) GetSegmentId() string { + if x != nil { + return x.SegmentId + } + return "" +} + +func (x *FlushSegmentCompactionInfo) GetFilePaths() map[string]*FilePaths { + if x != nil { + return x.FilePaths + } + return nil +} + +type FlushCollectionCompactionRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + TenantId string `protobuf:"bytes,1,opt,name=tenant_id,json=tenantId,proto3" json:"tenant_id,omitempty"` + CollectionId string `protobuf:"bytes,2,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + LogPosition int64 `protobuf:"varint,3,opt,name=log_position,json=logPosition,proto3" json:"log_position,omitempty"` + CollectionVersion int32 `protobuf:"varint,4,opt,name=collection_version,json=collectionVersion,proto3" json:"collection_version,omitempty"` + SegmentCompactionInfo []*FlushSegmentCompactionInfo `protobuf:"bytes,5,rep,name=segment_compaction_info,json=segmentCompactionInfo,proto3" json:"segment_compaction_info,omitempty"` +} + +func (x *FlushCollectionCompactionRequest) Reset() { + *x = FlushCollectionCompactionRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[31] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *FlushCollectionCompactionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FlushCollectionCompactionRequest) ProtoMessage() {} + +func (x *FlushCollectionCompactionRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[31] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FlushCollectionCompactionRequest.ProtoReflect.Descriptor instead. +func (*FlushCollectionCompactionRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{31} +} + +func (x *FlushCollectionCompactionRequest) GetTenantId() string { + if x != nil { + return x.TenantId + } + return "" +} + +func (x *FlushCollectionCompactionRequest) GetCollectionId() string { + if x != nil { + return x.CollectionId + } + return "" +} + +func (x *FlushCollectionCompactionRequest) GetLogPosition() int64 { + if x != nil { + return x.LogPosition + } + return 0 +} + +func (x *FlushCollectionCompactionRequest) GetCollectionVersion() int32 { + if x != nil { + return x.CollectionVersion + } + return 0 +} + +func (x *FlushCollectionCompactionRequest) GetSegmentCompactionInfo() []*FlushSegmentCompactionInfo { + if x != nil { + return x.SegmentCompactionInfo + } + return nil +} + +type FlushCollectionCompactionResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + CollectionId string `protobuf:"bytes,1,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + CollectionVersion int32 `protobuf:"varint,2,opt,name=collection_version,json=collectionVersion,proto3" json:"collection_version,omitempty"` + LastCompactionTime int64 `protobuf:"varint,3,opt,name=last_compaction_time,json=lastCompactionTime,proto3" json:"last_compaction_time,omitempty"` +} + +func (x *FlushCollectionCompactionResponse) Reset() { + *x = FlushCollectionCompactionResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[32] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *FlushCollectionCompactionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FlushCollectionCompactionResponse) ProtoMessage() {} + +func (x *FlushCollectionCompactionResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_coordinator_proto_msgTypes[32] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FlushCollectionCompactionResponse.ProtoReflect.Descriptor instead. +func (*FlushCollectionCompactionResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_coordinator_proto_rawDescGZIP(), []int{32} +} + +func (x *FlushCollectionCompactionResponse) GetCollectionId() string { + if x != nil { + return x.CollectionId + } + return "" +} + +func (x *FlushCollectionCompactionResponse) GetCollectionVersion() int32 { + if x != nil { + return x.CollectionVersion + } + return 0 +} + +func (x *FlushCollectionCompactionResponse) GetLastCompactionTime() int64 { + if x != nil { + return x.LastCompactionTime + } + return 0 +} + +var File_chromadb_proto_coordinator_proto protoreflect.FileDescriptor + +var file_chromadb_proto_coordinator_proto_rawDesc = []byte{ + 0x0a, 0x20, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x64, 0x62, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x2f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x12, 0x06, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x1a, 0x1b, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x64, 0x62, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x1b, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x53, 0x0a, 0x15, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x44, 0x61, + 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, + 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, + 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, + 0x65, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x22, 0x40, 0x0a, 0x16, 0x43, 0x72, 0x65, + 0x61, 0x74, 0x65, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, + 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x40, 0x0a, 0x12, 0x47, + 0x65, 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, + 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x22, 0x6b, 0x0a, + 0x13, 0x47, 0x65, 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x73, 0x70, + 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x2c, 0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, + 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, + 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x02, 0x20, 0x01, + 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, + 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x29, 0x0a, 0x13, 0x43, 0x72, + 0x65, 0x61, 0x74, 0x65, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x3e, 0x0a, 0x14, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x54, + 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, + 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, + 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x26, 0x0a, 0x10, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, + 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, + 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x63, 0x0a, + 0x11, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, + 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x54, 0x65, 0x6e, 0x61, + 0x6e, 0x74, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, + 0x61, 0x74, 0x75, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, + 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, + 0x75, 0x73, 0x22, 0x41, 0x0a, 0x14, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, + 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x29, 0x0a, 0x07, 0x73, 0x65, + 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0f, 0x2e, 0x63, 0x68, + 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x07, 0x73, 0x65, + 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x22, 0x3f, 0x0a, 0x15, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, + 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, + 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, + 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, + 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x26, 0x0a, 0x14, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, + 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, + 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22, 0x3f, + 0x0a, 0x15, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, + 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, + 0xc1, 0x01, 0x0a, 0x12, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, + 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x13, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, + 0x28, 0x09, 0x48, 0x00, 0x52, 0x02, 0x69, 0x64, 0x88, 0x01, 0x01, 0x12, 0x17, 0x0a, 0x04, 0x74, + 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x04, 0x74, 0x79, 0x70, + 0x65, 0x88, 0x01, 0x01, 0x12, 0x2f, 0x0a, 0x05, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x18, 0x03, 0x20, + 0x01, 0x28, 0x0e, 0x32, 0x14, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, + 0x6d, 0x65, 0x6e, 0x74, 0x53, 0x63, 0x6f, 0x70, 0x65, 0x48, 0x02, 0x52, 0x05, 0x73, 0x63, 0x6f, + 0x70, 0x65, 0x88, 0x01, 0x01, 0x12, 0x23, 0x0a, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, + 0x69, 0x6f, 0x6e, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x48, 0x03, 0x52, 0x0a, 0x63, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x42, 0x05, 0x0a, 0x03, 0x5f, 0x69, + 0x64, 0x42, 0x07, 0x0a, 0x05, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x42, 0x08, 0x0a, 0x06, 0x5f, 0x73, + 0x63, 0x6f, 0x70, 0x65, 0x42, 0x0d, 0x0a, 0x0b, 0x5f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, + 0x69, 0x6f, 0x6e, 0x22, 0x6a, 0x0a, 0x13, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, + 0x74, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x2b, 0x0a, 0x08, 0x73, 0x65, + 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x0f, 0x2e, 0x63, + 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x08, 0x73, + 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, + 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, + 0xfc, 0x01, 0x0a, 0x14, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, + 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x20, 0x0a, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x48, 0x00, 0x52, 0x0a, + 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x2b, 0x0a, 0x10, 0x72, 0x65, + 0x73, 0x65, 0x74, 0x5f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x05, + 0x20, 0x01, 0x28, 0x08, 0x48, 0x00, 0x52, 0x0f, 0x72, 0x65, 0x73, 0x65, 0x74, 0x43, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x34, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, + 0x61, 0x74, 0x61, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, + 0x61, 0x48, 0x01, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x12, 0x27, 0x0a, + 0x0e, 0x72, 0x65, 0x73, 0x65, 0x74, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, + 0x07, 0x20, 0x01, 0x28, 0x08, 0x48, 0x01, 0x52, 0x0d, 0x72, 0x65, 0x73, 0x65, 0x74, 0x4d, 0x65, + 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x42, 0x13, 0x0a, 0x11, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x42, 0x11, 0x0a, 0x0f, 0x6d, + 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x22, 0x3f, + 0x0a, 0x15, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, + 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, + 0xa3, 0x02, 0x0a, 0x17, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, + 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, + 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, + 0x37, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x03, 0x20, 0x01, 0x28, + 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, + 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x48, 0x00, 0x52, 0x08, 0x6d, 0x65, 0x74, + 0x61, 0x64, 0x61, 0x74, 0x61, 0x88, 0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, + 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x48, 0x01, 0x52, 0x09, 0x64, + 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x27, 0x0a, 0x0d, 0x67, + 0x65, 0x74, 0x5f, 0x6f, 0x72, 0x5f, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, 0x18, 0x05, 0x20, 0x01, + 0x28, 0x08, 0x48, 0x02, 0x52, 0x0b, 0x67, 0x65, 0x74, 0x4f, 0x72, 0x43, 0x72, 0x65, 0x61, 0x74, + 0x65, 0x88, 0x01, 0x01, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x06, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1a, 0x0a, 0x08, + 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, + 0x64, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x42, 0x0b, 0x0a, 0x09, 0x5f, 0x6d, 0x65, 0x74, + 0x61, 0x64, 0x61, 0x74, 0x61, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, + 0x69, 0x6f, 0x6e, 0x42, 0x10, 0x0a, 0x0e, 0x5f, 0x67, 0x65, 0x74, 0x5f, 0x6f, 0x72, 0x5f, 0x63, + 0x72, 0x65, 0x61, 0x74, 0x65, 0x22, 0x90, 0x01, 0x0a, 0x18, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, + 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x12, 0x32, 0x0a, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x0a, 0x63, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x18, 0x0a, 0x07, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, + 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, 0x64, + 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, + 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, + 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x5d, 0x0a, 0x17, 0x44, 0x65, 0x6c, 0x65, + 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, + 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x02, 0x69, 0x64, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x02, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x64, + 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, + 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x22, 0x42, 0x0a, 0x18, 0x44, 0x65, 0x6c, 0x65, 0x74, + 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, + 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x89, 0x01, 0x0a, 0x15, + 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x13, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, + 0x09, 0x48, 0x00, 0x52, 0x02, 0x69, 0x64, 0x88, 0x01, 0x01, 0x12, 0x17, 0x0a, 0x04, 0x6e, 0x61, + 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, + 0x88, 0x01, 0x01, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x18, 0x04, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x06, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x64, + 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, + 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x42, 0x05, 0x0a, 0x03, 0x5f, 0x69, 0x64, 0x42, 0x07, + 0x0a, 0x05, 0x5f, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x76, 0x0a, 0x16, 0x47, 0x65, 0x74, 0x43, 0x6f, + 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, + 0x65, 0x12, 0x34, 0x0a, 0x0b, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, + 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x0b, 0x63, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x12, 0x26, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, + 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, + 0xee, 0x01, 0x0a, 0x17, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, + 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x17, 0x0a, 0x04, 0x6e, + 0x61, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x48, 0x01, 0x52, 0x04, 0x6e, 0x61, 0x6d, + 0x65, 0x88, 0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, + 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x48, 0x02, 0x52, 0x09, 0x64, 0x69, 0x6d, 0x65, 0x6e, + 0x73, 0x69, 0x6f, 0x6e, 0x88, 0x01, 0x01, 0x12, 0x34, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, + 0x61, 0x74, 0x61, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, + 0x61, 0x48, 0x00, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x12, 0x27, 0x0a, + 0x0e, 0x72, 0x65, 0x73, 0x65, 0x74, 0x5f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, + 0x06, 0x20, 0x01, 0x28, 0x08, 0x48, 0x00, 0x52, 0x0d, 0x72, 0x65, 0x73, 0x65, 0x74, 0x4d, 0x65, + 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x42, 0x11, 0x0a, 0x0f, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, + 0x74, 0x61, 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x42, 0x07, 0x0a, 0x05, 0x5f, 0x6e, 0x61, + 0x6d, 0x65, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, + 0x22, 0x42, 0x0a, 0x18, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, + 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, + 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, + 0x61, 0x74, 0x75, 0x73, 0x22, 0x6f, 0x0a, 0x0c, 0x4e, 0x6f, 0x74, 0x69, 0x66, 0x69, 0x63, 0x61, + 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x03, + 0x52, 0x02, 0x69, 0x64, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, + 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, + 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x16, 0x0a, + 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73, + 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x3c, 0x0a, 0x12, 0x52, 0x65, 0x73, 0x65, 0x74, 0x53, 0x74, + 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x73, + 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x68, + 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x06, 0x73, 0x74, 0x61, + 0x74, 0x75, 0x73, 0x22, 0x44, 0x0a, 0x25, 0x47, 0x65, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, + 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x46, 0x6f, 0x72, 0x54, + 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1b, 0x0a, 0x09, + 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, + 0x08, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x49, 0x64, 0x22, 0x69, 0x0a, 0x18, 0x54, 0x65, 0x6e, + 0x61, 0x6e, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, + 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x12, 0x1b, 0x0a, 0x09, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x5f, + 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, + 0x49, 0x64, 0x12, 0x30, 0x0a, 0x14, 0x6c, 0x61, 0x73, 0x74, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x61, + 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x74, 0x69, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, + 0x52, 0x12, 0x6c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x54, 0x69, 0x6d, 0x65, 0x22, 0x89, 0x01, 0x0a, 0x26, 0x47, 0x65, 0x74, 0x4c, 0x61, 0x73, 0x74, + 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x46, 0x6f, + 0x72, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, + 0x5f, 0x0a, 0x1b, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x5f, 0x6c, 0x61, 0x73, 0x74, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x74, 0x69, 0x6d, 0x65, 0x18, 0x01, + 0x20, 0x03, 0x28, 0x0b, 0x32, 0x20, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x54, 0x65, + 0x6e, 0x61, 0x6e, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, + 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x52, 0x18, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x4c, 0x61, + 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, + 0x22, 0x88, 0x01, 0x0a, 0x25, 0x53, 0x65, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, + 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x46, 0x6f, 0x72, 0x54, 0x65, 0x6e, + 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x5f, 0x0a, 0x1b, 0x74, 0x65, + 0x6e, 0x61, 0x6e, 0x74, 0x5f, 0x6c, 0x61, 0x73, 0x74, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x61, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x74, 0x69, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, + 0x20, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x4c, + 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, + 0x65, 0x52, 0x18, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, + 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x22, 0xde, 0x01, 0x0a, 0x1a, + 0x46, 0x6c, 0x75, 0x73, 0x68, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6d, 0x70, + 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, + 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, + 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x12, 0x50, 0x0a, 0x0a, 0x66, 0x69, 0x6c, + 0x65, 0x5f, 0x70, 0x61, 0x74, 0x68, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x31, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x46, 0x6c, 0x75, 0x73, 0x68, 0x53, 0x65, 0x67, 0x6d, + 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, + 0x6f, 0x2e, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, + 0x52, 0x09, 0x66, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x1a, 0x4f, 0x0a, 0x0e, 0x46, + 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, + 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, + 0x27, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, + 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x61, 0x74, 0x68, + 0x73, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0x92, 0x02, 0x0a, + 0x20, 0x46, 0x6c, 0x75, 0x73, 0x68, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x1b, 0x0a, 0x09, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x74, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x49, 0x64, 0x12, 0x23, + 0x0a, 0x0d, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, + 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, + 0x6e, 0x49, 0x64, 0x12, 0x21, 0x0a, 0x0c, 0x6c, 0x6f, 0x67, 0x5f, 0x70, 0x6f, 0x73, 0x69, 0x74, + 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0b, 0x6c, 0x6f, 0x67, 0x50, 0x6f, + 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x2d, 0x0a, 0x12, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x04, 0x20, 0x01, + 0x28, 0x05, 0x52, 0x11, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x56, 0x65, + 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x5a, 0x0a, 0x17, 0x73, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, + 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x6e, 0x66, 0x6f, + 0x18, 0x05, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x22, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x46, 0x6c, 0x75, 0x73, 0x68, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6d, 0x70, + 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x15, 0x73, 0x65, 0x67, 0x6d, + 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, + 0x6f, 0x22, 0xa9, 0x01, 0x0a, 0x21, 0x46, 0x6c, 0x75, 0x73, 0x68, 0x43, 0x6f, 0x6c, 0x6c, 0x65, + 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6c, 0x6c, 0x65, + 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, + 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x2d, 0x0a, 0x12, + 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x76, 0x65, 0x72, 0x73, 0x69, + 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x11, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x56, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x30, 0x0a, 0x14, 0x6c, + 0x61, 0x73, 0x74, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x74, + 0x69, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x12, 0x6c, 0x61, 0x73, 0x74, 0x43, + 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x32, 0xf4, 0x0a, + 0x0a, 0x05, 0x53, 0x79, 0x73, 0x44, 0x42, 0x12, 0x51, 0x0a, 0x0e, 0x43, 0x72, 0x65, 0x61, 0x74, + 0x65, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x12, 0x1d, 0x2e, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, + 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, + 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x47, 0x65, + 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x12, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, + 0x65, 0x74, 0x44, 0x61, 0x74, 0x61, 0x62, 0x61, 0x73, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x22, 0x00, 0x12, 0x4b, 0x0a, 0x0c, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x54, 0x65, + 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, + 0x65, 0x61, 0x74, 0x65, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x1c, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, + 0x65, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, + 0x00, 0x12, 0x42, 0x0a, 0x09, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x18, + 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, + 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x19, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x2e, 0x47, 0x65, 0x74, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4e, 0x0a, 0x0d, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, + 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x1c, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, + 0x65, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4e, 0x0a, 0x0d, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, + 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x12, 0x1c, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, + 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x44, 0x65, + 0x6c, 0x65, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, 0x6d, + 0x65, 0x6e, 0x74, 0x73, 0x12, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, + 0x74, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x1b, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x53, 0x65, 0x67, + 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, + 0x4e, 0x0a, 0x0d, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, + 0x12, 0x1c, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, + 0x53, 0x65, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, + 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x53, 0x65, + 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, + 0x57, 0x0a, 0x10, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, + 0x69, 0x6f, 0x6e, 0x12, 0x1f, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, 0x65, + 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x72, + 0x65, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, + 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x57, 0x0a, 0x10, 0x44, 0x65, 0x6c, 0x65, + 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x1f, 0x2e, 0x63, + 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x43, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, + 0x00, 0x12, 0x51, 0x0a, 0x0e, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, + 0x6f, 0x6e, 0x73, 0x12, 0x1d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, + 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x43, + 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x22, 0x00, 0x12, 0x57, 0x0a, 0x10, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, + 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x1f, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, + 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, + 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x42, 0x0a, + 0x0a, 0x52, 0x65, 0x73, 0x65, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x12, 0x16, 0x2e, 0x67, 0x6f, + 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, + 0x70, 0x74, 0x79, 0x1a, 0x1a, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x52, 0x65, 0x73, + 0x65, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, + 0x00, 0x12, 0x81, 0x01, 0x0a, 0x1e, 0x47, 0x65, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, + 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x46, 0x6f, 0x72, 0x54, 0x65, + 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x2d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, + 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, + 0x69, 0x6d, 0x65, 0x46, 0x6f, 0x72, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, + 0x65, 0x73, 0x74, 0x1a, 0x2e, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, + 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, + 0x6d, 0x65, 0x46, 0x6f, 0x72, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x69, 0x0a, 0x1e, 0x53, 0x65, 0x74, 0x4c, 0x61, 0x73, 0x74, + 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x46, 0x6f, + 0x72, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x12, 0x2d, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x53, 0x65, 0x74, 0x4c, 0x61, 0x73, 0x74, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, + 0x6f, 0x6e, 0x54, 0x69, 0x6d, 0x65, 0x46, 0x6f, 0x72, 0x54, 0x65, 0x6e, 0x61, 0x6e, 0x74, 0x52, + 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x22, 0x00, + 0x12, 0x72, 0x0a, 0x19, 0x46, 0x6c, 0x75, 0x73, 0x68, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, + 0x69, 0x6f, 0x6e, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x28, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x46, 0x6c, 0x75, 0x73, 0x68, 0x43, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x29, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x46, 0x6c, 0x75, 0x73, 0x68, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x22, 0x00, 0x42, 0x3a, 0x5a, 0x38, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, + 0x6f, 0x6d, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2d, 0x63, 0x6f, 0x72, 0x65, 0x2f, 0x63, + 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2f, 0x67, 0x6f, 0x2f, 0x70, 0x6b, 0x67, 0x2f, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x2f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x70, 0x62, + 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, +} + +var ( + file_chromadb_proto_coordinator_proto_rawDescOnce sync.Once + file_chromadb_proto_coordinator_proto_rawDescData = file_chromadb_proto_coordinator_proto_rawDesc +) + +func file_chromadb_proto_coordinator_proto_rawDescGZIP() []byte { + file_chromadb_proto_coordinator_proto_rawDescOnce.Do(func() { + file_chromadb_proto_coordinator_proto_rawDescData = protoimpl.X.CompressGZIP(file_chromadb_proto_coordinator_proto_rawDescData) + }) + return file_chromadb_proto_coordinator_proto_rawDescData +} + +var file_chromadb_proto_coordinator_proto_msgTypes = make([]protoimpl.MessageInfo, 34) +var file_chromadb_proto_coordinator_proto_goTypes = []interface{}{ + (*CreateDatabaseRequest)(nil), // 0: chroma.CreateDatabaseRequest + (*CreateDatabaseResponse)(nil), // 1: chroma.CreateDatabaseResponse + (*GetDatabaseRequest)(nil), // 2: chroma.GetDatabaseRequest + (*GetDatabaseResponse)(nil), // 3: chroma.GetDatabaseResponse + (*CreateTenantRequest)(nil), // 4: chroma.CreateTenantRequest + (*CreateTenantResponse)(nil), // 5: chroma.CreateTenantResponse + (*GetTenantRequest)(nil), // 6: chroma.GetTenantRequest + (*GetTenantResponse)(nil), // 7: chroma.GetTenantResponse + (*CreateSegmentRequest)(nil), // 8: chroma.CreateSegmentRequest + (*CreateSegmentResponse)(nil), // 9: chroma.CreateSegmentResponse + (*DeleteSegmentRequest)(nil), // 10: chroma.DeleteSegmentRequest + (*DeleteSegmentResponse)(nil), // 11: chroma.DeleteSegmentResponse + (*GetSegmentsRequest)(nil), // 12: chroma.GetSegmentsRequest + (*GetSegmentsResponse)(nil), // 13: chroma.GetSegmentsResponse + (*UpdateSegmentRequest)(nil), // 14: chroma.UpdateSegmentRequest + (*UpdateSegmentResponse)(nil), // 15: chroma.UpdateSegmentResponse + (*CreateCollectionRequest)(nil), // 16: chroma.CreateCollectionRequest + (*CreateCollectionResponse)(nil), // 17: chroma.CreateCollectionResponse + (*DeleteCollectionRequest)(nil), // 18: chroma.DeleteCollectionRequest + (*DeleteCollectionResponse)(nil), // 19: chroma.DeleteCollectionResponse + (*GetCollectionsRequest)(nil), // 20: chroma.GetCollectionsRequest + (*GetCollectionsResponse)(nil), // 21: chroma.GetCollectionsResponse + (*UpdateCollectionRequest)(nil), // 22: chroma.UpdateCollectionRequest + (*UpdateCollectionResponse)(nil), // 23: chroma.UpdateCollectionResponse + (*Notification)(nil), // 24: chroma.Notification + (*ResetStateResponse)(nil), // 25: chroma.ResetStateResponse + (*GetLastCompactionTimeForTenantRequest)(nil), // 26: chroma.GetLastCompactionTimeForTenantRequest + (*TenantLastCompactionTime)(nil), // 27: chroma.TenantLastCompactionTime + (*GetLastCompactionTimeForTenantResponse)(nil), // 28: chroma.GetLastCompactionTimeForTenantResponse + (*SetLastCompactionTimeForTenantRequest)(nil), // 29: chroma.SetLastCompactionTimeForTenantRequest + (*FlushSegmentCompactionInfo)(nil), // 30: chroma.FlushSegmentCompactionInfo + (*FlushCollectionCompactionRequest)(nil), // 31: chroma.FlushCollectionCompactionRequest + (*FlushCollectionCompactionResponse)(nil), // 32: chroma.FlushCollectionCompactionResponse + nil, // 33: chroma.FlushSegmentCompactionInfo.FilePathsEntry + (*Status)(nil), // 34: chroma.Status + (*Database)(nil), // 35: chroma.Database + (*Tenant)(nil), // 36: chroma.Tenant + (*Segment)(nil), // 37: chroma.Segment + (SegmentScope)(0), // 38: chroma.SegmentScope + (*UpdateMetadata)(nil), // 39: chroma.UpdateMetadata + (*Collection)(nil), // 40: chroma.Collection + (*FilePaths)(nil), // 41: chroma.FilePaths + (*emptypb.Empty)(nil), // 42: google.protobuf.Empty +} +var file_chromadb_proto_coordinator_proto_depIdxs = []int32{ + 34, // 0: chroma.CreateDatabaseResponse.status:type_name -> chroma.Status + 35, // 1: chroma.GetDatabaseResponse.database:type_name -> chroma.Database + 34, // 2: chroma.GetDatabaseResponse.status:type_name -> chroma.Status + 34, // 3: chroma.CreateTenantResponse.status:type_name -> chroma.Status + 36, // 4: chroma.GetTenantResponse.tenant:type_name -> chroma.Tenant + 34, // 5: chroma.GetTenantResponse.status:type_name -> chroma.Status + 37, // 6: chroma.CreateSegmentRequest.segment:type_name -> chroma.Segment + 34, // 7: chroma.CreateSegmentResponse.status:type_name -> chroma.Status + 34, // 8: chroma.DeleteSegmentResponse.status:type_name -> chroma.Status + 38, // 9: chroma.GetSegmentsRequest.scope:type_name -> chroma.SegmentScope + 37, // 10: chroma.GetSegmentsResponse.segments:type_name -> chroma.Segment + 34, // 11: chroma.GetSegmentsResponse.status:type_name -> chroma.Status + 39, // 12: chroma.UpdateSegmentRequest.metadata:type_name -> chroma.UpdateMetadata + 34, // 13: chroma.UpdateSegmentResponse.status:type_name -> chroma.Status + 39, // 14: chroma.CreateCollectionRequest.metadata:type_name -> chroma.UpdateMetadata + 40, // 15: chroma.CreateCollectionResponse.collection:type_name -> chroma.Collection + 34, // 16: chroma.CreateCollectionResponse.status:type_name -> chroma.Status + 34, // 17: chroma.DeleteCollectionResponse.status:type_name -> chroma.Status + 40, // 18: chroma.GetCollectionsResponse.collections:type_name -> chroma.Collection + 34, // 19: chroma.GetCollectionsResponse.status:type_name -> chroma.Status + 39, // 20: chroma.UpdateCollectionRequest.metadata:type_name -> chroma.UpdateMetadata + 34, // 21: chroma.UpdateCollectionResponse.status:type_name -> chroma.Status + 34, // 22: chroma.ResetStateResponse.status:type_name -> chroma.Status + 27, // 23: chroma.GetLastCompactionTimeForTenantResponse.tenant_last_compaction_time:type_name -> chroma.TenantLastCompactionTime + 27, // 24: chroma.SetLastCompactionTimeForTenantRequest.tenant_last_compaction_time:type_name -> chroma.TenantLastCompactionTime + 33, // 25: chroma.FlushSegmentCompactionInfo.file_paths:type_name -> chroma.FlushSegmentCompactionInfo.FilePathsEntry + 30, // 26: chroma.FlushCollectionCompactionRequest.segment_compaction_info:type_name -> chroma.FlushSegmentCompactionInfo + 41, // 27: chroma.FlushSegmentCompactionInfo.FilePathsEntry.value:type_name -> chroma.FilePaths + 0, // 28: chroma.SysDB.CreateDatabase:input_type -> chroma.CreateDatabaseRequest + 2, // 29: chroma.SysDB.GetDatabase:input_type -> chroma.GetDatabaseRequest + 4, // 30: chroma.SysDB.CreateTenant:input_type -> chroma.CreateTenantRequest + 6, // 31: chroma.SysDB.GetTenant:input_type -> chroma.GetTenantRequest + 8, // 32: chroma.SysDB.CreateSegment:input_type -> chroma.CreateSegmentRequest + 10, // 33: chroma.SysDB.DeleteSegment:input_type -> chroma.DeleteSegmentRequest + 12, // 34: chroma.SysDB.GetSegments:input_type -> chroma.GetSegmentsRequest + 14, // 35: chroma.SysDB.UpdateSegment:input_type -> chroma.UpdateSegmentRequest + 16, // 36: chroma.SysDB.CreateCollection:input_type -> chroma.CreateCollectionRequest + 18, // 37: chroma.SysDB.DeleteCollection:input_type -> chroma.DeleteCollectionRequest + 20, // 38: chroma.SysDB.GetCollections:input_type -> chroma.GetCollectionsRequest + 22, // 39: chroma.SysDB.UpdateCollection:input_type -> chroma.UpdateCollectionRequest + 42, // 40: chroma.SysDB.ResetState:input_type -> google.protobuf.Empty + 26, // 41: chroma.SysDB.GetLastCompactionTimeForTenant:input_type -> chroma.GetLastCompactionTimeForTenantRequest + 29, // 42: chroma.SysDB.SetLastCompactionTimeForTenant:input_type -> chroma.SetLastCompactionTimeForTenantRequest + 31, // 43: chroma.SysDB.FlushCollectionCompaction:input_type -> chroma.FlushCollectionCompactionRequest + 1, // 44: chroma.SysDB.CreateDatabase:output_type -> chroma.CreateDatabaseResponse + 3, // 45: chroma.SysDB.GetDatabase:output_type -> chroma.GetDatabaseResponse + 5, // 46: chroma.SysDB.CreateTenant:output_type -> chroma.CreateTenantResponse + 7, // 47: chroma.SysDB.GetTenant:output_type -> chroma.GetTenantResponse + 9, // 48: chroma.SysDB.CreateSegment:output_type -> chroma.CreateSegmentResponse + 11, // 49: chroma.SysDB.DeleteSegment:output_type -> chroma.DeleteSegmentResponse + 13, // 50: chroma.SysDB.GetSegments:output_type -> chroma.GetSegmentsResponse + 15, // 51: chroma.SysDB.UpdateSegment:output_type -> chroma.UpdateSegmentResponse + 17, // 52: chroma.SysDB.CreateCollection:output_type -> chroma.CreateCollectionResponse + 19, // 53: chroma.SysDB.DeleteCollection:output_type -> chroma.DeleteCollectionResponse + 21, // 54: chroma.SysDB.GetCollections:output_type -> chroma.GetCollectionsResponse + 23, // 55: chroma.SysDB.UpdateCollection:output_type -> chroma.UpdateCollectionResponse + 25, // 56: chroma.SysDB.ResetState:output_type -> chroma.ResetStateResponse + 28, // 57: chroma.SysDB.GetLastCompactionTimeForTenant:output_type -> chroma.GetLastCompactionTimeForTenantResponse + 42, // 58: chroma.SysDB.SetLastCompactionTimeForTenant:output_type -> google.protobuf.Empty + 32, // 59: chroma.SysDB.FlushCollectionCompaction:output_type -> chroma.FlushCollectionCompactionResponse + 44, // [44:60] is the sub-list for method output_type + 28, // [28:44] is the sub-list for method input_type + 28, // [28:28] is the sub-list for extension type_name + 28, // [28:28] is the sub-list for extension extendee + 0, // [0:28] is the sub-list for field type_name +} + +func init() { file_chromadb_proto_coordinator_proto_init() } +func file_chromadb_proto_coordinator_proto_init() { + if File_chromadb_proto_coordinator_proto != nil { + return + } + file_chromadb_proto_chroma_proto_init() + if !protoimpl.UnsafeEnabled { + file_chromadb_proto_coordinator_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateDatabaseRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateDatabaseResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetDatabaseRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetDatabaseResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateTenantRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateTenantResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetTenantRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetTenantResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[8].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateSegmentRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[9].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateSegmentResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[10].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*DeleteSegmentRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[11].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*DeleteSegmentResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[12].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetSegmentsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[13].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetSegmentsResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[14].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*UpdateSegmentRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[15].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*UpdateSegmentResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[16].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateCollectionRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[17].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CreateCollectionResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[18].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*DeleteCollectionRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[19].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*DeleteCollectionResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[20].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetCollectionsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[21].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetCollectionsResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[22].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*UpdateCollectionRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[23].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*UpdateCollectionResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[24].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*Notification); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[25].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*ResetStateResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[26].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetLastCompactionTimeForTenantRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[27].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*TenantLastCompactionTime); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[28].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetLastCompactionTimeForTenantResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[29].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*SetLastCompactionTimeForTenantRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[30].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*FlushSegmentCompactionInfo); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[31].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*FlushCollectionCompactionRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_coordinator_proto_msgTypes[32].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*FlushCollectionCompactionResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + } + file_chromadb_proto_coordinator_proto_msgTypes[12].OneofWrappers = []interface{}{} + file_chromadb_proto_coordinator_proto_msgTypes[14].OneofWrappers = []interface{}{ + (*UpdateSegmentRequest_Collection)(nil), + (*UpdateSegmentRequest_ResetCollection)(nil), + (*UpdateSegmentRequest_Metadata)(nil), + (*UpdateSegmentRequest_ResetMetadata)(nil), + } + file_chromadb_proto_coordinator_proto_msgTypes[16].OneofWrappers = []interface{}{} + file_chromadb_proto_coordinator_proto_msgTypes[20].OneofWrappers = []interface{}{} + file_chromadb_proto_coordinator_proto_msgTypes[22].OneofWrappers = []interface{}{ + (*UpdateCollectionRequest_Metadata)(nil), + (*UpdateCollectionRequest_ResetMetadata)(nil), + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: file_chromadb_proto_coordinator_proto_rawDesc, + NumEnums: 0, + NumMessages: 34, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_chromadb_proto_coordinator_proto_goTypes, + DependencyIndexes: file_chromadb_proto_coordinator_proto_depIdxs, + MessageInfos: file_chromadb_proto_coordinator_proto_msgTypes, + }.Build() + File_chromadb_proto_coordinator_proto = out.File + file_chromadb_proto_coordinator_proto_rawDesc = nil + file_chromadb_proto_coordinator_proto_goTypes = nil + file_chromadb_proto_coordinator_proto_depIdxs = nil +} diff --git a/go/coordinator/internal/proto/coordinatorpb/coordinator_grpc.pb.go b/go/pkg/proto/coordinatorpb/coordinator_grpc.pb.go similarity index 66% rename from go/coordinator/internal/proto/coordinatorpb/coordinator_grpc.pb.go rename to go/pkg/proto/coordinatorpb/coordinator_grpc.pb.go index ed123f9f3a6..1306dbc1793 100644 --- a/go/coordinator/internal/proto/coordinatorpb/coordinator_grpc.pb.go +++ b/go/pkg/proto/coordinatorpb/coordinator_grpc.pb.go @@ -20,38 +20,44 @@ import ( const _ = grpc.SupportPackageIsVersion7 const ( - SysDB_CreateDatabase_FullMethodName = "/chroma.SysDB/CreateDatabase" - SysDB_GetDatabase_FullMethodName = "/chroma.SysDB/GetDatabase" - SysDB_CreateTenant_FullMethodName = "/chroma.SysDB/CreateTenant" - SysDB_GetTenant_FullMethodName = "/chroma.SysDB/GetTenant" - SysDB_CreateSegment_FullMethodName = "/chroma.SysDB/CreateSegment" - SysDB_DeleteSegment_FullMethodName = "/chroma.SysDB/DeleteSegment" - SysDB_GetSegments_FullMethodName = "/chroma.SysDB/GetSegments" - SysDB_UpdateSegment_FullMethodName = "/chroma.SysDB/UpdateSegment" - SysDB_CreateCollection_FullMethodName = "/chroma.SysDB/CreateCollection" - SysDB_DeleteCollection_FullMethodName = "/chroma.SysDB/DeleteCollection" - SysDB_GetCollections_FullMethodName = "/chroma.SysDB/GetCollections" - SysDB_UpdateCollection_FullMethodName = "/chroma.SysDB/UpdateCollection" - SysDB_ResetState_FullMethodName = "/chroma.SysDB/ResetState" + SysDB_CreateDatabase_FullMethodName = "/chroma.SysDB/CreateDatabase" + SysDB_GetDatabase_FullMethodName = "/chroma.SysDB/GetDatabase" + SysDB_CreateTenant_FullMethodName = "/chroma.SysDB/CreateTenant" + SysDB_GetTenant_FullMethodName = "/chroma.SysDB/GetTenant" + SysDB_CreateSegment_FullMethodName = "/chroma.SysDB/CreateSegment" + SysDB_DeleteSegment_FullMethodName = "/chroma.SysDB/DeleteSegment" + SysDB_GetSegments_FullMethodName = "/chroma.SysDB/GetSegments" + SysDB_UpdateSegment_FullMethodName = "/chroma.SysDB/UpdateSegment" + SysDB_CreateCollection_FullMethodName = "/chroma.SysDB/CreateCollection" + SysDB_DeleteCollection_FullMethodName = "/chroma.SysDB/DeleteCollection" + SysDB_GetCollections_FullMethodName = "/chroma.SysDB/GetCollections" + SysDB_UpdateCollection_FullMethodName = "/chroma.SysDB/UpdateCollection" + SysDB_ResetState_FullMethodName = "/chroma.SysDB/ResetState" + SysDB_GetLastCompactionTimeForTenant_FullMethodName = "/chroma.SysDB/GetLastCompactionTimeForTenant" + SysDB_SetLastCompactionTimeForTenant_FullMethodName = "/chroma.SysDB/SetLastCompactionTimeForTenant" + SysDB_FlushCollectionCompaction_FullMethodName = "/chroma.SysDB/FlushCollectionCompaction" ) // SysDBClient is the client API for SysDB service. // // For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. type SysDBClient interface { - CreateDatabase(ctx context.Context, in *CreateDatabaseRequest, opts ...grpc.CallOption) (*ChromaResponse, error) + CreateDatabase(ctx context.Context, in *CreateDatabaseRequest, opts ...grpc.CallOption) (*CreateDatabaseResponse, error) GetDatabase(ctx context.Context, in *GetDatabaseRequest, opts ...grpc.CallOption) (*GetDatabaseResponse, error) - CreateTenant(ctx context.Context, in *CreateTenantRequest, opts ...grpc.CallOption) (*ChromaResponse, error) + CreateTenant(ctx context.Context, in *CreateTenantRequest, opts ...grpc.CallOption) (*CreateTenantResponse, error) GetTenant(ctx context.Context, in *GetTenantRequest, opts ...grpc.CallOption) (*GetTenantResponse, error) - CreateSegment(ctx context.Context, in *CreateSegmentRequest, opts ...grpc.CallOption) (*ChromaResponse, error) - DeleteSegment(ctx context.Context, in *DeleteSegmentRequest, opts ...grpc.CallOption) (*ChromaResponse, error) + CreateSegment(ctx context.Context, in *CreateSegmentRequest, opts ...grpc.CallOption) (*CreateSegmentResponse, error) + DeleteSegment(ctx context.Context, in *DeleteSegmentRequest, opts ...grpc.CallOption) (*DeleteSegmentResponse, error) GetSegments(ctx context.Context, in *GetSegmentsRequest, opts ...grpc.CallOption) (*GetSegmentsResponse, error) - UpdateSegment(ctx context.Context, in *UpdateSegmentRequest, opts ...grpc.CallOption) (*ChromaResponse, error) + UpdateSegment(ctx context.Context, in *UpdateSegmentRequest, opts ...grpc.CallOption) (*UpdateSegmentResponse, error) CreateCollection(ctx context.Context, in *CreateCollectionRequest, opts ...grpc.CallOption) (*CreateCollectionResponse, error) - DeleteCollection(ctx context.Context, in *DeleteCollectionRequest, opts ...grpc.CallOption) (*ChromaResponse, error) + DeleteCollection(ctx context.Context, in *DeleteCollectionRequest, opts ...grpc.CallOption) (*DeleteCollectionResponse, error) GetCollections(ctx context.Context, in *GetCollectionsRequest, opts ...grpc.CallOption) (*GetCollectionsResponse, error) - UpdateCollection(ctx context.Context, in *UpdateCollectionRequest, opts ...grpc.CallOption) (*ChromaResponse, error) - ResetState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*ChromaResponse, error) + UpdateCollection(ctx context.Context, in *UpdateCollectionRequest, opts ...grpc.CallOption) (*UpdateCollectionResponse, error) + ResetState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*ResetStateResponse, error) + GetLastCompactionTimeForTenant(ctx context.Context, in *GetLastCompactionTimeForTenantRequest, opts ...grpc.CallOption) (*GetLastCompactionTimeForTenantResponse, error) + SetLastCompactionTimeForTenant(ctx context.Context, in *SetLastCompactionTimeForTenantRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) + FlushCollectionCompaction(ctx context.Context, in *FlushCollectionCompactionRequest, opts ...grpc.CallOption) (*FlushCollectionCompactionResponse, error) } type sysDBClient struct { @@ -62,8 +68,8 @@ func NewSysDBClient(cc grpc.ClientConnInterface) SysDBClient { return &sysDBClient{cc} } -func (c *sysDBClient) CreateDatabase(ctx context.Context, in *CreateDatabaseRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) CreateDatabase(ctx context.Context, in *CreateDatabaseRequest, opts ...grpc.CallOption) (*CreateDatabaseResponse, error) { + out := new(CreateDatabaseResponse) err := c.cc.Invoke(ctx, SysDB_CreateDatabase_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -80,8 +86,8 @@ func (c *sysDBClient) GetDatabase(ctx context.Context, in *GetDatabaseRequest, o return out, nil } -func (c *sysDBClient) CreateTenant(ctx context.Context, in *CreateTenantRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) CreateTenant(ctx context.Context, in *CreateTenantRequest, opts ...grpc.CallOption) (*CreateTenantResponse, error) { + out := new(CreateTenantResponse) err := c.cc.Invoke(ctx, SysDB_CreateTenant_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -98,8 +104,8 @@ func (c *sysDBClient) GetTenant(ctx context.Context, in *GetTenantRequest, opts return out, nil } -func (c *sysDBClient) CreateSegment(ctx context.Context, in *CreateSegmentRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) CreateSegment(ctx context.Context, in *CreateSegmentRequest, opts ...grpc.CallOption) (*CreateSegmentResponse, error) { + out := new(CreateSegmentResponse) err := c.cc.Invoke(ctx, SysDB_CreateSegment_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -107,8 +113,8 @@ func (c *sysDBClient) CreateSegment(ctx context.Context, in *CreateSegmentReques return out, nil } -func (c *sysDBClient) DeleteSegment(ctx context.Context, in *DeleteSegmentRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) DeleteSegment(ctx context.Context, in *DeleteSegmentRequest, opts ...grpc.CallOption) (*DeleteSegmentResponse, error) { + out := new(DeleteSegmentResponse) err := c.cc.Invoke(ctx, SysDB_DeleteSegment_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -125,8 +131,8 @@ func (c *sysDBClient) GetSegments(ctx context.Context, in *GetSegmentsRequest, o return out, nil } -func (c *sysDBClient) UpdateSegment(ctx context.Context, in *UpdateSegmentRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) UpdateSegment(ctx context.Context, in *UpdateSegmentRequest, opts ...grpc.CallOption) (*UpdateSegmentResponse, error) { + out := new(UpdateSegmentResponse) err := c.cc.Invoke(ctx, SysDB_UpdateSegment_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -143,8 +149,8 @@ func (c *sysDBClient) CreateCollection(ctx context.Context, in *CreateCollection return out, nil } -func (c *sysDBClient) DeleteCollection(ctx context.Context, in *DeleteCollectionRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) DeleteCollection(ctx context.Context, in *DeleteCollectionRequest, opts ...grpc.CallOption) (*DeleteCollectionResponse, error) { + out := new(DeleteCollectionResponse) err := c.cc.Invoke(ctx, SysDB_DeleteCollection_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -161,8 +167,8 @@ func (c *sysDBClient) GetCollections(ctx context.Context, in *GetCollectionsRequ return out, nil } -func (c *sysDBClient) UpdateCollection(ctx context.Context, in *UpdateCollectionRequest, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) UpdateCollection(ctx context.Context, in *UpdateCollectionRequest, opts ...grpc.CallOption) (*UpdateCollectionResponse, error) { + out := new(UpdateCollectionResponse) err := c.cc.Invoke(ctx, SysDB_UpdateCollection_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -170,8 +176,8 @@ func (c *sysDBClient) UpdateCollection(ctx context.Context, in *UpdateCollection return out, nil } -func (c *sysDBClient) ResetState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*ChromaResponse, error) { - out := new(ChromaResponse) +func (c *sysDBClient) ResetState(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*ResetStateResponse, error) { + out := new(ResetStateResponse) err := c.cc.Invoke(ctx, SysDB_ResetState_FullMethodName, in, out, opts...) if err != nil { return nil, err @@ -179,23 +185,53 @@ func (c *sysDBClient) ResetState(ctx context.Context, in *emptypb.Empty, opts .. return out, nil } +func (c *sysDBClient) GetLastCompactionTimeForTenant(ctx context.Context, in *GetLastCompactionTimeForTenantRequest, opts ...grpc.CallOption) (*GetLastCompactionTimeForTenantResponse, error) { + out := new(GetLastCompactionTimeForTenantResponse) + err := c.cc.Invoke(ctx, SysDB_GetLastCompactionTimeForTenant_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *sysDBClient) SetLastCompactionTimeForTenant(ctx context.Context, in *SetLastCompactionTimeForTenantRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) { + out := new(emptypb.Empty) + err := c.cc.Invoke(ctx, SysDB_SetLastCompactionTimeForTenant_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *sysDBClient) FlushCollectionCompaction(ctx context.Context, in *FlushCollectionCompactionRequest, opts ...grpc.CallOption) (*FlushCollectionCompactionResponse, error) { + out := new(FlushCollectionCompactionResponse) + err := c.cc.Invoke(ctx, SysDB_FlushCollectionCompaction_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + // SysDBServer is the server API for SysDB service. // All implementations must embed UnimplementedSysDBServer // for forward compatibility type SysDBServer interface { - CreateDatabase(context.Context, *CreateDatabaseRequest) (*ChromaResponse, error) + CreateDatabase(context.Context, *CreateDatabaseRequest) (*CreateDatabaseResponse, error) GetDatabase(context.Context, *GetDatabaseRequest) (*GetDatabaseResponse, error) - CreateTenant(context.Context, *CreateTenantRequest) (*ChromaResponse, error) + CreateTenant(context.Context, *CreateTenantRequest) (*CreateTenantResponse, error) GetTenant(context.Context, *GetTenantRequest) (*GetTenantResponse, error) - CreateSegment(context.Context, *CreateSegmentRequest) (*ChromaResponse, error) - DeleteSegment(context.Context, *DeleteSegmentRequest) (*ChromaResponse, error) + CreateSegment(context.Context, *CreateSegmentRequest) (*CreateSegmentResponse, error) + DeleteSegment(context.Context, *DeleteSegmentRequest) (*DeleteSegmentResponse, error) GetSegments(context.Context, *GetSegmentsRequest) (*GetSegmentsResponse, error) - UpdateSegment(context.Context, *UpdateSegmentRequest) (*ChromaResponse, error) + UpdateSegment(context.Context, *UpdateSegmentRequest) (*UpdateSegmentResponse, error) CreateCollection(context.Context, *CreateCollectionRequest) (*CreateCollectionResponse, error) - DeleteCollection(context.Context, *DeleteCollectionRequest) (*ChromaResponse, error) + DeleteCollection(context.Context, *DeleteCollectionRequest) (*DeleteCollectionResponse, error) GetCollections(context.Context, *GetCollectionsRequest) (*GetCollectionsResponse, error) - UpdateCollection(context.Context, *UpdateCollectionRequest) (*ChromaResponse, error) - ResetState(context.Context, *emptypb.Empty) (*ChromaResponse, error) + UpdateCollection(context.Context, *UpdateCollectionRequest) (*UpdateCollectionResponse, error) + ResetState(context.Context, *emptypb.Empty) (*ResetStateResponse, error) + GetLastCompactionTimeForTenant(context.Context, *GetLastCompactionTimeForTenantRequest) (*GetLastCompactionTimeForTenantResponse, error) + SetLastCompactionTimeForTenant(context.Context, *SetLastCompactionTimeForTenantRequest) (*emptypb.Empty, error) + FlushCollectionCompaction(context.Context, *FlushCollectionCompactionRequest) (*FlushCollectionCompactionResponse, error) mustEmbedUnimplementedSysDBServer() } @@ -203,45 +239,54 @@ type SysDBServer interface { type UnimplementedSysDBServer struct { } -func (UnimplementedSysDBServer) CreateDatabase(context.Context, *CreateDatabaseRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) CreateDatabase(context.Context, *CreateDatabaseRequest) (*CreateDatabaseResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method CreateDatabase not implemented") } func (UnimplementedSysDBServer) GetDatabase(context.Context, *GetDatabaseRequest) (*GetDatabaseResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method GetDatabase not implemented") } -func (UnimplementedSysDBServer) CreateTenant(context.Context, *CreateTenantRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) CreateTenant(context.Context, *CreateTenantRequest) (*CreateTenantResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method CreateTenant not implemented") } func (UnimplementedSysDBServer) GetTenant(context.Context, *GetTenantRequest) (*GetTenantResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method GetTenant not implemented") } -func (UnimplementedSysDBServer) CreateSegment(context.Context, *CreateSegmentRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) CreateSegment(context.Context, *CreateSegmentRequest) (*CreateSegmentResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method CreateSegment not implemented") } -func (UnimplementedSysDBServer) DeleteSegment(context.Context, *DeleteSegmentRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) DeleteSegment(context.Context, *DeleteSegmentRequest) (*DeleteSegmentResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method DeleteSegment not implemented") } func (UnimplementedSysDBServer) GetSegments(context.Context, *GetSegmentsRequest) (*GetSegmentsResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method GetSegments not implemented") } -func (UnimplementedSysDBServer) UpdateSegment(context.Context, *UpdateSegmentRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) UpdateSegment(context.Context, *UpdateSegmentRequest) (*UpdateSegmentResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method UpdateSegment not implemented") } func (UnimplementedSysDBServer) CreateCollection(context.Context, *CreateCollectionRequest) (*CreateCollectionResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method CreateCollection not implemented") } -func (UnimplementedSysDBServer) DeleteCollection(context.Context, *DeleteCollectionRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) DeleteCollection(context.Context, *DeleteCollectionRequest) (*DeleteCollectionResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method DeleteCollection not implemented") } func (UnimplementedSysDBServer) GetCollections(context.Context, *GetCollectionsRequest) (*GetCollectionsResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method GetCollections not implemented") } -func (UnimplementedSysDBServer) UpdateCollection(context.Context, *UpdateCollectionRequest) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) UpdateCollection(context.Context, *UpdateCollectionRequest) (*UpdateCollectionResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method UpdateCollection not implemented") } -func (UnimplementedSysDBServer) ResetState(context.Context, *emptypb.Empty) (*ChromaResponse, error) { +func (UnimplementedSysDBServer) ResetState(context.Context, *emptypb.Empty) (*ResetStateResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method ResetState not implemented") } +func (UnimplementedSysDBServer) GetLastCompactionTimeForTenant(context.Context, *GetLastCompactionTimeForTenantRequest) (*GetLastCompactionTimeForTenantResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method GetLastCompactionTimeForTenant not implemented") +} +func (UnimplementedSysDBServer) SetLastCompactionTimeForTenant(context.Context, *SetLastCompactionTimeForTenantRequest) (*emptypb.Empty, error) { + return nil, status.Errorf(codes.Unimplemented, "method SetLastCompactionTimeForTenant not implemented") +} +func (UnimplementedSysDBServer) FlushCollectionCompaction(context.Context, *FlushCollectionCompactionRequest) (*FlushCollectionCompactionResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method FlushCollectionCompaction not implemented") +} func (UnimplementedSysDBServer) mustEmbedUnimplementedSysDBServer() {} // UnsafeSysDBServer may be embedded to opt out of forward compatibility for this service. @@ -489,6 +534,60 @@ func _SysDB_ResetState_Handler(srv interface{}, ctx context.Context, dec func(in return interceptor(ctx, in, info, handler) } +func _SysDB_GetLastCompactionTimeForTenant_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(GetLastCompactionTimeForTenantRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(SysDBServer).GetLastCompactionTimeForTenant(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: SysDB_GetLastCompactionTimeForTenant_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(SysDBServer).GetLastCompactionTimeForTenant(ctx, req.(*GetLastCompactionTimeForTenantRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _SysDB_SetLastCompactionTimeForTenant_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(SetLastCompactionTimeForTenantRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(SysDBServer).SetLastCompactionTimeForTenant(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: SysDB_SetLastCompactionTimeForTenant_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(SysDBServer).SetLastCompactionTimeForTenant(ctx, req.(*SetLastCompactionTimeForTenantRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _SysDB_FlushCollectionCompaction_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(FlushCollectionCompactionRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(SysDBServer).FlushCollectionCompaction(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: SysDB_FlushCollectionCompaction_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(SysDBServer).FlushCollectionCompaction(ctx, req.(*FlushCollectionCompactionRequest)) + } + return interceptor(ctx, in, info, handler) +} + // SysDB_ServiceDesc is the grpc.ServiceDesc for SysDB service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -548,6 +647,18 @@ var SysDB_ServiceDesc = grpc.ServiceDesc{ MethodName: "ResetState", Handler: _SysDB_ResetState_Handler, }, + { + MethodName: "GetLastCompactionTimeForTenant", + Handler: _SysDB_GetLastCompactionTimeForTenant_Handler, + }, + { + MethodName: "SetLastCompactionTimeForTenant", + Handler: _SysDB_SetLastCompactionTimeForTenant_Handler, + }, + { + MethodName: "FlushCollectionCompaction", + Handler: _SysDB_FlushCollectionCompaction_Handler, + }, }, Streams: []grpc.StreamDesc{}, Metadata: "chromadb/proto/coordinator.proto", diff --git a/go/pkg/proto/logservicepb/logservice.pb.go b/go/pkg/proto/logservicepb/logservice.pb.go new file mode 100644 index 00000000000..434aa10f898 --- /dev/null +++ b/go/pkg/proto/logservicepb/logservice.pb.go @@ -0,0 +1,692 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.33.0 +// protoc v4.23.4 +// source: chromadb/proto/logservice.proto + +package logservicepb + +import ( + coordinatorpb "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb" + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type PushLogsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + CollectionId string `protobuf:"bytes,1,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + Records []*coordinatorpb.OperationRecord `protobuf:"bytes,2,rep,name=records,proto3" json:"records,omitempty"` +} + +func (x *PushLogsRequest) Reset() { + *x = PushLogsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PushLogsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PushLogsRequest) ProtoMessage() {} + +func (x *PushLogsRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[0] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PushLogsRequest.ProtoReflect.Descriptor instead. +func (*PushLogsRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{0} +} + +func (x *PushLogsRequest) GetCollectionId() string { + if x != nil { + return x.CollectionId + } + return "" +} + +func (x *PushLogsRequest) GetRecords() []*coordinatorpb.OperationRecord { + if x != nil { + return x.Records + } + return nil +} + +type PushLogsResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + RecordCount int32 `protobuf:"varint,1,opt,name=record_count,json=recordCount,proto3" json:"record_count,omitempty"` +} + +func (x *PushLogsResponse) Reset() { + *x = PushLogsResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PushLogsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PushLogsResponse) ProtoMessage() {} + +func (x *PushLogsResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[1] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PushLogsResponse.ProtoReflect.Descriptor instead. +func (*PushLogsResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{1} +} + +func (x *PushLogsResponse) GetRecordCount() int32 { + if x != nil { + return x.RecordCount + } + return 0 +} + +type PullLogsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + CollectionId string `protobuf:"bytes,1,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + StartFromOffset int64 `protobuf:"varint,2,opt,name=start_from_offset,json=startFromOffset,proto3" json:"start_from_offset,omitempty"` + BatchSize int32 `protobuf:"varint,3,opt,name=batch_size,json=batchSize,proto3" json:"batch_size,omitempty"` + EndTimestamp int64 `protobuf:"varint,4,opt,name=end_timestamp,json=endTimestamp,proto3" json:"end_timestamp,omitempty"` +} + +func (x *PullLogsRequest) Reset() { + *x = PullLogsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PullLogsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PullLogsRequest) ProtoMessage() {} + +func (x *PullLogsRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[2] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PullLogsRequest.ProtoReflect.Descriptor instead. +func (*PullLogsRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{2} +} + +func (x *PullLogsRequest) GetCollectionId() string { + if x != nil { + return x.CollectionId + } + return "" +} + +func (x *PullLogsRequest) GetStartFromOffset() int64 { + if x != nil { + return x.StartFromOffset + } + return 0 +} + +func (x *PullLogsRequest) GetBatchSize() int32 { + if x != nil { + return x.BatchSize + } + return 0 +} + +func (x *PullLogsRequest) GetEndTimestamp() int64 { + if x != nil { + return x.EndTimestamp + } + return 0 +} + +// Represents an operation from the log +type LogRecord struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + LogOffset int64 `protobuf:"varint,1,opt,name=log_offset,json=logOffset,proto3" json:"log_offset,omitempty"` + Record *coordinatorpb.OperationRecord `protobuf:"bytes,2,opt,name=record,proto3" json:"record,omitempty"` +} + +func (x *LogRecord) Reset() { + *x = LogRecord{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *LogRecord) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*LogRecord) ProtoMessage() {} + +func (x *LogRecord) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[3] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use LogRecord.ProtoReflect.Descriptor instead. +func (*LogRecord) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{3} +} + +func (x *LogRecord) GetLogOffset() int64 { + if x != nil { + return x.LogOffset + } + return 0 +} + +func (x *LogRecord) GetRecord() *coordinatorpb.OperationRecord { + if x != nil { + return x.Record + } + return nil +} + +type PullLogsResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Records []*LogRecord `protobuf:"bytes,1,rep,name=records,proto3" json:"records,omitempty"` +} + +func (x *PullLogsResponse) Reset() { + *x = PullLogsResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PullLogsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PullLogsResponse) ProtoMessage() {} + +func (x *PullLogsResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[4] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PullLogsResponse.ProtoReflect.Descriptor instead. +func (*PullLogsResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{4} +} + +func (x *PullLogsResponse) GetRecords() []*LogRecord { + if x != nil { + return x.Records + } + return nil +} + +type CollectionInfo struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + CollectionId string `protobuf:"bytes,1,opt,name=collection_id,json=collectionId,proto3" json:"collection_id,omitempty"` + // The log offset of the first log entry of the collection that needs to be compacted + FirstLogOffset int64 `protobuf:"varint,2,opt,name=first_log_offset,json=firstLogOffset,proto3" json:"first_log_offset,omitempty"` + // The timestamp of the first log entry of the collection that needs to be compacted + FirstLogTs int64 `protobuf:"varint,3,opt,name=first_log_ts,json=firstLogTs,proto3" json:"first_log_ts,omitempty"` +} + +func (x *CollectionInfo) Reset() { + *x = CollectionInfo{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CollectionInfo) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CollectionInfo) ProtoMessage() {} + +func (x *CollectionInfo) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[5] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CollectionInfo.ProtoReflect.Descriptor instead. +func (*CollectionInfo) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{5} +} + +func (x *CollectionInfo) GetCollectionId() string { + if x != nil { + return x.CollectionId + } + return "" +} + +func (x *CollectionInfo) GetFirstLogOffset() int64 { + if x != nil { + return x.FirstLogOffset + } + return 0 +} + +func (x *CollectionInfo) GetFirstLogTs() int64 { + if x != nil { + return x.FirstLogTs + } + return 0 +} + +type GetAllCollectionInfoToCompactRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *GetAllCollectionInfoToCompactRequest) Reset() { + *x = GetAllCollectionInfoToCompactRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetAllCollectionInfoToCompactRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetAllCollectionInfoToCompactRequest) ProtoMessage() {} + +func (x *GetAllCollectionInfoToCompactRequest) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[6] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetAllCollectionInfoToCompactRequest.ProtoReflect.Descriptor instead. +func (*GetAllCollectionInfoToCompactRequest) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{6} +} + +type GetAllCollectionInfoToCompactResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + AllCollectionInfo []*CollectionInfo `protobuf:"bytes,1,rep,name=all_collection_info,json=allCollectionInfo,proto3" json:"all_collection_info,omitempty"` +} + +func (x *GetAllCollectionInfoToCompactResponse) Reset() { + *x = GetAllCollectionInfoToCompactResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_chromadb_proto_logservice_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetAllCollectionInfoToCompactResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetAllCollectionInfoToCompactResponse) ProtoMessage() {} + +func (x *GetAllCollectionInfoToCompactResponse) ProtoReflect() protoreflect.Message { + mi := &file_chromadb_proto_logservice_proto_msgTypes[7] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetAllCollectionInfoToCompactResponse.ProtoReflect.Descriptor instead. +func (*GetAllCollectionInfoToCompactResponse) Descriptor() ([]byte, []int) { + return file_chromadb_proto_logservice_proto_rawDescGZIP(), []int{7} +} + +func (x *GetAllCollectionInfoToCompactResponse) GetAllCollectionInfo() []*CollectionInfo { + if x != nil { + return x.AllCollectionInfo + } + return nil +} + +var File_chromadb_proto_logservice_proto protoreflect.FileDescriptor + +var file_chromadb_proto_logservice_proto_rawDesc = []byte{ + 0x0a, 0x1f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x64, 0x62, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x2f, 0x6c, 0x6f, 0x67, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x12, 0x06, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x1a, 0x1b, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x64, 0x62, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x69, 0x0a, 0x0f, 0x50, 0x75, 0x73, 0x68, 0x4c, 0x6f, + 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, + 0x52, 0x0c, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x31, + 0x0a, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28, 0x0b, 0x32, + 0x17, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x4f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, + 0x6f, 0x6e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x52, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x73, 0x22, 0x35, 0x0a, 0x10, 0x50, 0x75, 0x73, 0x68, 0x4c, 0x6f, 0x67, 0x73, 0x52, 0x65, 0x73, + 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x21, 0x0a, 0x0c, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x5f, + 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0b, 0x72, 0x65, 0x63, + 0x6f, 0x72, 0x64, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x22, 0xa6, 0x01, 0x0a, 0x0f, 0x50, 0x75, 0x6c, + 0x6c, 0x4c, 0x6f, 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x23, 0x0a, 0x0d, + 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, + 0x64, 0x12, 0x2a, 0x0a, 0x11, 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f, 0x66, 0x72, 0x6f, 0x6d, 0x5f, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0f, 0x73, 0x74, + 0x61, 0x72, 0x74, 0x46, 0x72, 0x6f, 0x6d, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x12, 0x1d, 0x0a, + 0x0a, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, + 0x05, 0x52, 0x09, 0x62, 0x61, 0x74, 0x63, 0x68, 0x53, 0x69, 0x7a, 0x65, 0x12, 0x23, 0x0a, 0x0d, + 0x65, 0x6e, 0x64, 0x5f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, + 0x01, 0x28, 0x03, 0x52, 0x0c, 0x65, 0x6e, 0x64, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, + 0x70, 0x22, 0x5b, 0x0a, 0x09, 0x4c, 0x6f, 0x67, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x12, 0x1d, + 0x0a, 0x0a, 0x6c, 0x6f, 0x67, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x18, 0x01, 0x20, 0x01, + 0x28, 0x03, 0x52, 0x09, 0x6c, 0x6f, 0x67, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x12, 0x2f, 0x0a, + 0x06, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x17, 0x2e, + 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x4f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, + 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x52, 0x06, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x22, 0x3f, + 0x0a, 0x10, 0x50, 0x75, 0x6c, 0x6c, 0x4c, 0x6f, 0x67, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x12, 0x2b, 0x0a, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x18, 0x01, 0x20, + 0x03, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x4c, 0x6f, 0x67, + 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x52, 0x07, 0x72, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x22, + 0x81, 0x01, 0x0a, 0x0e, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, + 0x66, 0x6f, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6f, 0x6c, 0x6c, 0x65, + 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x28, 0x0a, 0x10, 0x66, 0x69, 0x72, 0x73, 0x74, + 0x5f, 0x6c, 0x6f, 0x67, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, + 0x03, 0x52, 0x0e, 0x66, 0x69, 0x72, 0x73, 0x74, 0x4c, 0x6f, 0x67, 0x4f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x12, 0x20, 0x0a, 0x0c, 0x66, 0x69, 0x72, 0x73, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x5f, 0x74, + 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0a, 0x66, 0x69, 0x72, 0x73, 0x74, 0x4c, 0x6f, + 0x67, 0x54, 0x73, 0x22, 0x26, 0x0a, 0x24, 0x47, 0x65, 0x74, 0x41, 0x6c, 0x6c, 0x43, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x54, 0x6f, 0x43, 0x6f, 0x6d, + 0x70, 0x61, 0x63, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x6f, 0x0a, 0x25, 0x47, + 0x65, 0x74, 0x41, 0x6c, 0x6c, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, + 0x6e, 0x66, 0x6f, 0x54, 0x6f, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x52, 0x65, 0x73, 0x70, + 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x46, 0x0a, 0x13, 0x61, 0x6c, 0x6c, 0x5f, 0x63, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x6e, 0x66, 0x6f, 0x18, 0x01, 0x20, 0x03, 0x28, + 0x0b, 0x32, 0x16, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x43, 0x6f, 0x6c, 0x6c, 0x65, + 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x11, 0x61, 0x6c, 0x6c, 0x43, 0x6f, + 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x32, 0x8e, 0x02, 0x0a, + 0x0a, 0x4c, 0x6f, 0x67, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x3f, 0x0a, 0x08, 0x50, + 0x75, 0x73, 0x68, 0x4c, 0x6f, 0x67, 0x73, 0x12, 0x17, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, + 0x2e, 0x50, 0x75, 0x73, 0x68, 0x4c, 0x6f, 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x18, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x50, 0x75, 0x73, 0x68, 0x4c, 0x6f, + 0x67, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x3f, 0x0a, 0x08, + 0x50, 0x75, 0x6c, 0x6c, 0x4c, 0x6f, 0x67, 0x73, 0x12, 0x17, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, + 0x61, 0x2e, 0x50, 0x75, 0x6c, 0x6c, 0x4c, 0x6f, 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x18, 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x50, 0x75, 0x6c, 0x6c, 0x4c, + 0x6f, 0x67, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x7e, 0x0a, + 0x1d, 0x47, 0x65, 0x74, 0x41, 0x6c, 0x6c, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, + 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x54, 0x6f, 0x43, 0x6f, 0x6d, 0x70, 0x61, 0x63, 0x74, 0x12, 0x2c, + 0x2e, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x6c, 0x6c, 0x43, 0x6f, + 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x54, 0x6f, 0x43, 0x6f, + 0x6d, 0x70, 0x61, 0x63, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x2d, 0x2e, 0x63, + 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x6c, 0x6c, 0x43, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x54, 0x6f, 0x43, 0x6f, 0x6d, 0x70, + 0x61, 0x63, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x39, 0x5a, + 0x37, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x68, 0x72, 0x6f, + 0x6d, 0x61, 0x2d, 0x63, 0x6f, 0x72, 0x65, 0x2f, 0x63, 0x68, 0x72, 0x6f, 0x6d, 0x61, 0x2f, 0x67, + 0x6f, 0x2f, 0x70, 0x6b, 0x67, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x6c, 0x6f, 0x67, 0x73, + 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, +} + +var ( + file_chromadb_proto_logservice_proto_rawDescOnce sync.Once + file_chromadb_proto_logservice_proto_rawDescData = file_chromadb_proto_logservice_proto_rawDesc +) + +func file_chromadb_proto_logservice_proto_rawDescGZIP() []byte { + file_chromadb_proto_logservice_proto_rawDescOnce.Do(func() { + file_chromadb_proto_logservice_proto_rawDescData = protoimpl.X.CompressGZIP(file_chromadb_proto_logservice_proto_rawDescData) + }) + return file_chromadb_proto_logservice_proto_rawDescData +} + +var file_chromadb_proto_logservice_proto_msgTypes = make([]protoimpl.MessageInfo, 8) +var file_chromadb_proto_logservice_proto_goTypes = []interface{}{ + (*PushLogsRequest)(nil), // 0: chroma.PushLogsRequest + (*PushLogsResponse)(nil), // 1: chroma.PushLogsResponse + (*PullLogsRequest)(nil), // 2: chroma.PullLogsRequest + (*LogRecord)(nil), // 3: chroma.LogRecord + (*PullLogsResponse)(nil), // 4: chroma.PullLogsResponse + (*CollectionInfo)(nil), // 5: chroma.CollectionInfo + (*GetAllCollectionInfoToCompactRequest)(nil), // 6: chroma.GetAllCollectionInfoToCompactRequest + (*GetAllCollectionInfoToCompactResponse)(nil), // 7: chroma.GetAllCollectionInfoToCompactResponse + (*coordinatorpb.OperationRecord)(nil), // 8: chroma.OperationRecord +} +var file_chromadb_proto_logservice_proto_depIdxs = []int32{ + 8, // 0: chroma.PushLogsRequest.records:type_name -> chroma.OperationRecord + 8, // 1: chroma.LogRecord.record:type_name -> chroma.OperationRecord + 3, // 2: chroma.PullLogsResponse.records:type_name -> chroma.LogRecord + 5, // 3: chroma.GetAllCollectionInfoToCompactResponse.all_collection_info:type_name -> chroma.CollectionInfo + 0, // 4: chroma.LogService.PushLogs:input_type -> chroma.PushLogsRequest + 2, // 5: chroma.LogService.PullLogs:input_type -> chroma.PullLogsRequest + 6, // 6: chroma.LogService.GetAllCollectionInfoToCompact:input_type -> chroma.GetAllCollectionInfoToCompactRequest + 1, // 7: chroma.LogService.PushLogs:output_type -> chroma.PushLogsResponse + 4, // 8: chroma.LogService.PullLogs:output_type -> chroma.PullLogsResponse + 7, // 9: chroma.LogService.GetAllCollectionInfoToCompact:output_type -> chroma.GetAllCollectionInfoToCompactResponse + 7, // [7:10] is the sub-list for method output_type + 4, // [4:7] is the sub-list for method input_type + 4, // [4:4] is the sub-list for extension type_name + 4, // [4:4] is the sub-list for extension extendee + 0, // [0:4] is the sub-list for field type_name +} + +func init() { file_chromadb_proto_logservice_proto_init() } +func file_chromadb_proto_logservice_proto_init() { + if File_chromadb_proto_logservice_proto != nil { + return + } + if !protoimpl.UnsafeEnabled { + file_chromadb_proto_logservice_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*PushLogsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*PushLogsResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*PullLogsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*LogRecord); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*PullLogsResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*CollectionInfo); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetAllCollectionInfoToCompactRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_chromadb_proto_logservice_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*GetAllCollectionInfoToCompactResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: file_chromadb_proto_logservice_proto_rawDesc, + NumEnums: 0, + NumMessages: 8, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_chromadb_proto_logservice_proto_goTypes, + DependencyIndexes: file_chromadb_proto_logservice_proto_depIdxs, + MessageInfos: file_chromadb_proto_logservice_proto_msgTypes, + }.Build() + File_chromadb_proto_logservice_proto = out.File + file_chromadb_proto_logservice_proto_rawDesc = nil + file_chromadb_proto_logservice_proto_goTypes = nil + file_chromadb_proto_logservice_proto_depIdxs = nil +} diff --git a/go/pkg/proto/logservicepb/logservice_grpc.pb.go b/go/pkg/proto/logservicepb/logservice_grpc.pb.go new file mode 100644 index 00000000000..7b9895d172a --- /dev/null +++ b/go/pkg/proto/logservicepb/logservice_grpc.pb.go @@ -0,0 +1,183 @@ +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.3.0 +// - protoc v4.23.4 +// source: chromadb/proto/logservice.proto + +package logservicepb + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.32.0 or later. +const _ = grpc.SupportPackageIsVersion7 + +const ( + LogService_PushLogs_FullMethodName = "/chroma.LogService/PushLogs" + LogService_PullLogs_FullMethodName = "/chroma.LogService/PullLogs" + LogService_GetAllCollectionInfoToCompact_FullMethodName = "/chroma.LogService/GetAllCollectionInfoToCompact" +) + +// LogServiceClient is the client API for LogService service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type LogServiceClient interface { + PushLogs(ctx context.Context, in *PushLogsRequest, opts ...grpc.CallOption) (*PushLogsResponse, error) + PullLogs(ctx context.Context, in *PullLogsRequest, opts ...grpc.CallOption) (*PullLogsResponse, error) + GetAllCollectionInfoToCompact(ctx context.Context, in *GetAllCollectionInfoToCompactRequest, opts ...grpc.CallOption) (*GetAllCollectionInfoToCompactResponse, error) +} + +type logServiceClient struct { + cc grpc.ClientConnInterface +} + +func NewLogServiceClient(cc grpc.ClientConnInterface) LogServiceClient { + return &logServiceClient{cc} +} + +func (c *logServiceClient) PushLogs(ctx context.Context, in *PushLogsRequest, opts ...grpc.CallOption) (*PushLogsResponse, error) { + out := new(PushLogsResponse) + err := c.cc.Invoke(ctx, LogService_PushLogs_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *logServiceClient) PullLogs(ctx context.Context, in *PullLogsRequest, opts ...grpc.CallOption) (*PullLogsResponse, error) { + out := new(PullLogsResponse) + err := c.cc.Invoke(ctx, LogService_PullLogs_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *logServiceClient) GetAllCollectionInfoToCompact(ctx context.Context, in *GetAllCollectionInfoToCompactRequest, opts ...grpc.CallOption) (*GetAllCollectionInfoToCompactResponse, error) { + out := new(GetAllCollectionInfoToCompactResponse) + err := c.cc.Invoke(ctx, LogService_GetAllCollectionInfoToCompact_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + +// LogServiceServer is the server API for LogService service. +// All implementations must embed UnimplementedLogServiceServer +// for forward compatibility +type LogServiceServer interface { + PushLogs(context.Context, *PushLogsRequest) (*PushLogsResponse, error) + PullLogs(context.Context, *PullLogsRequest) (*PullLogsResponse, error) + GetAllCollectionInfoToCompact(context.Context, *GetAllCollectionInfoToCompactRequest) (*GetAllCollectionInfoToCompactResponse, error) + mustEmbedUnimplementedLogServiceServer() +} + +// UnimplementedLogServiceServer must be embedded to have forward compatible implementations. +type UnimplementedLogServiceServer struct { +} + +func (UnimplementedLogServiceServer) PushLogs(context.Context, *PushLogsRequest) (*PushLogsResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method PushLogs not implemented") +} +func (UnimplementedLogServiceServer) PullLogs(context.Context, *PullLogsRequest) (*PullLogsResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method PullLogs not implemented") +} +func (UnimplementedLogServiceServer) GetAllCollectionInfoToCompact(context.Context, *GetAllCollectionInfoToCompactRequest) (*GetAllCollectionInfoToCompactResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method GetAllCollectionInfoToCompact not implemented") +} +func (UnimplementedLogServiceServer) mustEmbedUnimplementedLogServiceServer() {} + +// UnsafeLogServiceServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to LogServiceServer will +// result in compilation errors. +type UnsafeLogServiceServer interface { + mustEmbedUnimplementedLogServiceServer() +} + +func RegisterLogServiceServer(s grpc.ServiceRegistrar, srv LogServiceServer) { + s.RegisterService(&LogService_ServiceDesc, srv) +} + +func _LogService_PushLogs_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PushLogsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(LogServiceServer).PushLogs(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: LogService_PushLogs_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(LogServiceServer).PushLogs(ctx, req.(*PushLogsRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _LogService_PullLogs_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PullLogsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(LogServiceServer).PullLogs(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: LogService_PullLogs_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(LogServiceServer).PullLogs(ctx, req.(*PullLogsRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _LogService_GetAllCollectionInfoToCompact_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(GetAllCollectionInfoToCompactRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(LogServiceServer).GetAllCollectionInfoToCompact(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: LogService_GetAllCollectionInfoToCompact_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(LogServiceServer).GetAllCollectionInfoToCompact(ctx, req.(*GetAllCollectionInfoToCompactRequest)) + } + return interceptor(ctx, in, info, handler) +} + +// LogService_ServiceDesc is the grpc.ServiceDesc for LogService service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var LogService_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "chroma.LogService", + HandlerType: (*LogServiceServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "PushLogs", + Handler: _LogService_PushLogs_Handler, + }, + { + MethodName: "PullLogs", + Handler: _LogService_PullLogs_Handler, + }, + { + MethodName: "GetAllCollectionInfoToCompact", + Handler: _LogService_GetAllCollectionInfoToCompact_Handler, + }, + }, + Streams: []grpc.StreamDesc{}, + Metadata: "chromadb/proto/logservice.proto", +} diff --git a/go/coordinator/internal/types/types.go b/go/pkg/types/types.go similarity index 100% rename from go/coordinator/internal/types/types.go rename to go/pkg/types/types.go diff --git a/go/coordinator/internal/utils/integration.go b/go/pkg/utils/integration.go similarity index 100% rename from go/coordinator/internal/utils/integration.go rename to go/pkg/utils/integration.go diff --git a/go/coordinator/internal/utils/kubernetes.go b/go/pkg/utils/kubernetes.go similarity index 100% rename from go/coordinator/internal/utils/kubernetes.go rename to go/pkg/utils/kubernetes.go diff --git a/go/coordinator/internal/utils/log.go b/go/pkg/utils/log.go similarity index 100% rename from go/coordinator/internal/utils/log.go rename to go/pkg/utils/log.go diff --git a/go/coordinator/internal/utils/pulsar_admin.go b/go/pkg/utils/pulsar_admin.go similarity index 100% rename from go/coordinator/internal/utils/pulsar_admin.go rename to go/pkg/utils/pulsar_admin.go diff --git a/go/coordinator/internal/utils/rendezvous_hash.go b/go/pkg/utils/rendezvous_hash.go similarity index 100% rename from go/coordinator/internal/utils/rendezvous_hash.go rename to go/pkg/utils/rendezvous_hash.go diff --git a/go/coordinator/internal/utils/rendezvous_hash_test.go b/go/pkg/utils/rendezvous_hash_test.go similarity index 100% rename from go/coordinator/internal/utils/rendezvous_hash_test.go rename to go/pkg/utils/rendezvous_hash_test.go diff --git a/go/coordinator/internal/utils/run.go b/go/pkg/utils/run.go similarity index 100% rename from go/coordinator/internal/utils/run.go rename to go/pkg/utils/run.go diff --git a/go/coordinator/internal/utils/signal.go b/go/pkg/utils/signal.go similarity index 100% rename from go/coordinator/internal/utils/signal.go rename to go/pkg/utils/signal.go diff --git a/go/shared/otel/main.go b/go/shared/otel/main.go new file mode 100644 index 00000000000..677ef344192 --- /dev/null +++ b/go/shared/otel/main.go @@ -0,0 +1,143 @@ +package otel + +import ( + "context" + "encoding/hex" + "fmt" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + otelCode "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.4.0" + "go.opentelemetry.io/otel/trace" + "google.golang.org/grpc" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +var tracer trace.Tracer + +func decodeTraceID(encodedSpanID string) (t trace.TraceID, err error) { + var spanBytes []byte + spanBytes, err = hex.DecodeString(encodedSpanID) + if err != nil { + err = fmt.Errorf("failed to decode spanID: %w", err) + return + } + copy(t[:], spanBytes) + return +} + +func decodeSpanID(encodedSpanID string) (s trace.SpanID, err error) { + var spanBytes []byte + spanBytes, err = hex.DecodeString(encodedSpanID) + if err != nil { + err = fmt.Errorf("failed to decode spanID: %w", err) + return + } + copy(s[:], spanBytes) + return +} + +// ServerGrpcInterceptor is a gRPC server interceptor for tracing and optional metadata-based context enhancement. +func ServerGrpcInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + // Init with a default tracer if not already set. (Unit test) + if tracer == nil { + tracer = otel.GetTracerProvider().Tracer("LOCAL") + } + // Attempt to retrieve metadata, but proceed normally if not present. + md, _ := metadata.FromIncomingContext(ctx) + + // Attempt to decode and apply trace and span IDs if present, without failing on their absence. + spanIdValue := decodeMetadataValue(md, "chroma-spanid") + traceIdValue := decodeMetadataValue(md, "chroma-traceid") + + var spanContext trace.SpanContext + if spanIdValue != "" && traceIdValue != "" { + if spanId, err := decodeSpanID(spanIdValue); err == nil { + if traceId, err := decodeTraceID(traceIdValue); err == nil { + spanContext = trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: traceId, + SpanID: spanId, + }) + // Only set the remote span context if both trace and span IDs are valid and decoded. + ctx = trace.ContextWithRemoteSpanContext(ctx, spanContext) + } + } + } + var span trace.Span + ctx, span = tracer.Start(ctx, "Request "+info.FullMethod) + defer span.End() + + // Calls the handler + h, err := handler(ctx, req) + if err != nil { + // Handle and log the error. + handleError(span, err) + return nil, err + } + + // Set the status to OK upon success. + span.SetStatus(otelCode.Ok, "ok") + span.SetAttributes(attribute.String("rpc.status_code", "ok")) + span.SetAttributes(attribute.String("rpc.method", info.FullMethod)) + + return h, nil +} + +// handleError logs and annotates the span with details of the encountered error. +func handleError(span trace.Span, err error) { + st, _ := status.FromError(err) + span.SetStatus(otelCode.Error, "error") + span.SetAttributes( + attribute.String("rpc.status_code", st.Code().String()), + attribute.String("rpc.message", st.Message()), + attribute.String("rpc.error", st.Err().Error()), + ) +} + +// decodeMetadataValue safely extracts a value from metadata, allowing for missing keys. +func decodeMetadataValue(md metadata.MD, key string) string { + values := md.Get(key) + if len(values) > 0 { + return values[0] + } + return "" +} + +type TracingConfig struct { + Endpoint string + Service string +} + +func InitTracing(ctx context.Context, config *TracingConfig) (err error) { + var exp *otlptrace.Exporter + exp, err = otlptrace.New( + ctx, + otlptracegrpc.NewClient( + otlptracegrpc.WithInsecure(), + otlptracegrpc.WithEndpoint(config.Endpoint), + otlptracegrpc.WithDialOption(grpc.WithBlock()), // Useful for waiting until the connection is up. + ), + ) + if err != nil { + return + } + + // Create a new tracer provider with a batch span processor and the OTLP exporter. + tp := sdktrace.NewTracerProvider( + sdktrace.WithBatcher(exp), + sdktrace.WithSampler(sdktrace.AlwaysSample()), + sdktrace.WithResource(resource.NewWithAttributes( + semconv.SchemaURL, + semconv.ServiceNameKey.String(config.Service), + )), + ) + + otel.SetTracerProvider(tp) + tracer = otel.Tracer(config.Service) + return +} diff --git a/idl/chromadb/proto/chroma.proto b/idl/chromadb/proto/chroma.proto index 5676c0efb74..1b3c8d2f311 100644 --- a/idl/chromadb/proto/chroma.proto +++ b/idl/chromadb/proto/chroma.proto @@ -2,19 +2,14 @@ syntax = "proto3"; package chroma; -option go_package = "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb"; +option go_package = "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"; message Status { string reason = 1; int32 code = 2; // TODO: What is the enum of this code? } -message ChromaResponse { - Status status = 1; -} - // Types here should mirror chromadb/types.py - enum Operation { ADD = 0; UPDATE = 1; @@ -38,25 +33,28 @@ enum SegmentScope { METADATA = 1; } +message FilePaths { + repeated string paths = 1; +} + message Segment { string id = 1; string type = 2; SegmentScope scope = 3; - optional string topic = 4; // TODO should channel <> segment binding exist here? - // If a segment has a collection, it implies that this segment implements the full - // collection and can be used to service queries (for it's given scope.) optional string collection = 5; optional UpdateMetadata metadata = 6; + map file_paths = 7; } message Collection { string id = 1; string name = 2; - string topic = 3; optional UpdateMetadata metadata = 4; optional int32 dimension = 5; string tenant = 6; string database = 7; + int64 log_position = 8; + int32 version = 9; } message Database { @@ -81,23 +79,21 @@ message UpdateMetadata { map metadata = 1; } -message SubmitEmbeddingRecord { +// Represents an operation the user submits +message OperationRecord { string id = 1; optional Vector vector = 2; optional UpdateMetadata metadata = 3; Operation operation = 4; - string collection_id = 5; } message VectorEmbeddingRecord { string id = 1; - bytes seq_id = 2; Vector vector = 3; // TODO: we need to rethink source of truth for vector dimensionality and encoding } message VectorQueryResult { string id = 1; - bytes seq_id = 2; float distance = 3; optional Vector vector = 4; } diff --git a/idl/chromadb/proto/coordinator.proto b/idl/chromadb/proto/coordinator.proto index 79abd73acf6..43668bdfad2 100644 --- a/idl/chromadb/proto/coordinator.proto +++ b/idl/chromadb/proto/coordinator.proto @@ -1,7 +1,7 @@ syntax = "proto3"; package chroma; -option go_package = "github.com/chroma/chroma-coordinator/internal/proto/coordinatorpb"; +option go_package = "github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"; import "chromadb/proto/chroma.proto"; import "google/protobuf/empty.proto"; @@ -12,6 +12,10 @@ message CreateDatabaseRequest { string tenant = 3; } +message CreateDatabaseResponse { + Status status = 1; +} + message GetDatabaseRequest { string name = 1; string tenant = 2; @@ -26,6 +30,10 @@ message CreateTenantRequest { string name = 2; // Names are globally unique } +message CreateTenantResponse { + Status status = 1; +} + message GetTenantRequest { string name = 1; } @@ -40,15 +48,22 @@ message CreateSegmentRequest { Segment segment = 1; } +message CreateSegmentResponse { + Status status = 1; +} + message DeleteSegmentRequest { string id = 1; } +message DeleteSegmentResponse { + Status status = 1; +} + message GetSegmentsRequest { optional string id = 1; optional string type = 2; optional SegmentScope scope = 3; - optional string topic = 4; optional string collection = 5; // Collection ID } @@ -60,10 +75,6 @@ message GetSegmentsResponse { message UpdateSegmentRequest { string id = 1; - oneof topic_update { - string topic = 2; - bool reset_topic = 3; - } oneof collection_update { string collection = 4; bool reset_collection = 5; @@ -74,6 +85,10 @@ message UpdateSegmentRequest { } } +message UpdateSegmentResponse { + Status status = 1; +} + message CreateCollectionRequest { string id = 1; string name = 2; @@ -96,10 +111,13 @@ message DeleteCollectionRequest { string database = 3; } +message DeleteCollectionResponse { + Status status = 1; +} + message GetCollectionsRequest { optional string id = 1; optional string name = 2; - optional string topic = 3; string tenant = 4; string database = 5; } @@ -111,7 +129,6 @@ message GetCollectionsResponse { message UpdateCollectionRequest { string id = 1; - optional string topic = 2; optional string name = 3; optional int32 dimension = 4; oneof metadata_update { @@ -120,6 +137,10 @@ message UpdateCollectionRequest { } } +message UpdateCollectionResponse { + Status status = 1; +} + message Notification { int64 id = 1; string collection_id = 2; @@ -127,18 +148,61 @@ message Notification { string status = 4; } +message ResetStateResponse { + Status status = 1; +} + +message GetLastCompactionTimeForTenantRequest { + repeated string tenant_id = 1; +} + +message TenantLastCompactionTime { + string tenant_id = 1; + int64 last_compaction_time = 2; +} + +message GetLastCompactionTimeForTenantResponse { + repeated TenantLastCompactionTime tenant_last_compaction_time = 1; +} + +message SetLastCompactionTimeForTenantRequest { + TenantLastCompactionTime tenant_last_compaction_time = 1; +} + +message FlushSegmentCompactionInfo { + string segment_id = 1; + map file_paths = 2; +} + +message FlushCollectionCompactionRequest { + string tenant_id = 1; + string collection_id = 2; + int64 log_position = 3; + int32 collection_version = 4; + repeated FlushSegmentCompactionInfo segment_compaction_info = 5; +} + +message FlushCollectionCompactionResponse { + string collection_id = 1; + int32 collection_version = 2; + int64 last_compaction_time = 3; +} + service SysDB { - rpc CreateDatabase(CreateDatabaseRequest) returns (ChromaResponse) {} + rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse) {} rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {} - rpc CreateTenant(CreateTenantRequest) returns (ChromaResponse) {} + rpc CreateTenant(CreateTenantRequest) returns (CreateTenantResponse) {} rpc GetTenant(GetTenantRequest) returns (GetTenantResponse) {} - rpc CreateSegment(CreateSegmentRequest) returns (ChromaResponse) {} - rpc DeleteSegment(DeleteSegmentRequest) returns (ChromaResponse) {} + rpc CreateSegment(CreateSegmentRequest) returns (CreateSegmentResponse) {} + rpc DeleteSegment(DeleteSegmentRequest) returns (DeleteSegmentResponse) {} rpc GetSegments(GetSegmentsRequest) returns (GetSegmentsResponse) {} - rpc UpdateSegment(UpdateSegmentRequest) returns (ChromaResponse) {} + rpc UpdateSegment(UpdateSegmentRequest) returns (UpdateSegmentResponse) {} rpc CreateCollection(CreateCollectionRequest) returns (CreateCollectionResponse) {} - rpc DeleteCollection(DeleteCollectionRequest) returns (ChromaResponse) {} + rpc DeleteCollection(DeleteCollectionRequest) returns (DeleteCollectionResponse) {} rpc GetCollections(GetCollectionsRequest) returns (GetCollectionsResponse) {} - rpc UpdateCollection(UpdateCollectionRequest) returns (ChromaResponse) {} - rpc ResetState(google.protobuf.Empty) returns (ChromaResponse) {} + rpc UpdateCollection(UpdateCollectionRequest) returns (UpdateCollectionResponse) {} + rpc ResetState(google.protobuf.Empty) returns (ResetStateResponse) {} + rpc GetLastCompactionTimeForTenant(GetLastCompactionTimeForTenantRequest) returns (GetLastCompactionTimeForTenantResponse) {} + rpc SetLastCompactionTimeForTenant(SetLastCompactionTimeForTenantRequest) returns (google.protobuf.Empty) {} + rpc FlushCollectionCompaction(FlushCollectionCompactionRequest) returns (FlushCollectionCompactionResponse) {} } diff --git a/idl/chromadb/proto/logservice.proto b/idl/chromadb/proto/logservice.proto new file mode 100644 index 00000000000..8c52a3165ce --- /dev/null +++ b/idl/chromadb/proto/logservice.proto @@ -0,0 +1,54 @@ +syntax = "proto3"; + +package chroma; +option go_package = "github.com/chroma-core/chroma/go/pkg/proto/logservicepb"; + +import "chromadb/proto/chroma.proto"; + +message PushLogsRequest { + string collection_id = 1; + repeated OperationRecord records = 2; +} + +message PushLogsResponse { + int32 record_count = 1; +} + +message PullLogsRequest { + string collection_id = 1; + int64 start_from_offset = 2; + int32 batch_size = 3; + int64 end_timestamp = 4; +} + +// Represents an operation from the log +message LogRecord { + int64 log_offset = 1; + OperationRecord record = 2; +} + +message PullLogsResponse { + repeated LogRecord records = 1; +} + +message CollectionInfo { + string collection_id = 1; + // The log offset of the first log entry of the collection that needs to be compacted + int64 first_log_offset = 2; + // The timestamp of the first log entry of the collection that needs to be compacted + int64 first_log_ts = 3; +} + +message GetAllCollectionInfoToCompactRequest { + // Empty +} + +message GetAllCollectionInfoToCompactResponse { + repeated CollectionInfo all_collection_info = 1; +} + +service LogService { + rpc PushLogs(PushLogsRequest) returns (PushLogsResponse) {} + rpc PullLogs(PullLogsRequest) returns (PullLogsResponse) {} + rpc GetAllCollectionInfoToCompact(GetAllCollectionInfoToCompactRequest) returns (GetAllCollectionInfoToCompactResponse) {} +} diff --git a/idl/makefile b/idl/makefile index 18cbc1977ba..22664e7b711 100644 --- a/idl/makefile +++ b/idl/makefile @@ -10,13 +10,14 @@ proto_python: proto_go: @echo "Generating gRPC code for golang..." @protoc \ - --go_out=../go/coordinator/internal/proto/coordinatorpb \ + --go_out=../go/pkg/proto/coordinatorpb \ --go_opt paths=source_relative \ --plugin protoc-gen-go="${GOPATH}/bin/protoc-gen-go" \ - --go-grpc_out=../go/coordinator/internal/proto/coordinatorpb \ + --go-grpc_out=../go/pkg/proto/coordinatorpb \ --go-grpc_opt paths=source_relative \ --plugin protoc-gen-go-grpc="${GOPATH}/bin/protoc-gen-go-grpc" \ chromadb/proto/*.proto - @mv ../go/coordinator/internal/proto/coordinatorpb/chromadb/proto/*.go ../go/coordinator/internal/proto/coordinatorpb/ - @rm -rf ../go/coordinator/internal/proto/coordinatorpb/chromadb + @mv ../go/pkg/proto/coordinatorpb/chromadb/proto/logservice*.go ../go/pkg/proto/logservicepb/ + @mv ../go/pkg/proto/coordinatorpb/chromadb/proto/*.go ../go/pkg/proto/coordinatorpb/ + @rm -rf ../go/pkg/proto/coordinatorpb/chromadb @echo "Done" diff --git a/k8s/cr/worker_memberlist_cr.yaml b/k8s/cr/worker_memberlist_cr.yaml deleted file mode 100644 index bc4df07f535..00000000000 --- a/k8s/cr/worker_memberlist_cr.yaml +++ /dev/null @@ -1,48 +0,0 @@ -# These kubernetes manifests are UNDER ACTIVE DEVELOPMENT and are not yet ready for production use. -# They will be used for the upcoming distributed version of chroma. They are not even ready -# for testing yet. Please do not use them unless you are working on the distributed version of chroma. - -# Create a memberlist called worker-memberlist -apiVersion: chroma.cluster/v1 -kind: MemberList -metadata: - name: worker-memberlist - namespace: chroma -spec: - members: - ---- - -apiVersion: rbac.authorization.k8s.io/v1 -kind: ClusterRole -metadata: - name: worker-memberlist-reader -rules: -- apiGroups: - - chroma.cluster - resources: - - memberlists - verbs: - - get - - list - - watch - # TODO: FIX THIS LEAKY PERMISSION - - create - - update - - patch - - delete - ---- - -apiVersion: rbac.authorization.k8s.io/v1 -kind: ClusterRoleBinding -metadata: - name: worker-memberlist-reader-binding -roleRef: - apiGroup: rbac.authorization.k8s.io - kind: ClusterRole - name: worker-memberlist-reader -subjects: -- kind: ServiceAccount - name: default - namespace: chroma diff --git a/k8s/deployment/kubernetes.yaml b/k8s/deployment/kubernetes.yaml deleted file mode 100644 index b1f9baabdd0..00000000000 --- a/k8s/deployment/kubernetes.yaml +++ /dev/null @@ -1,221 +0,0 @@ -# These kubernetes manifests are UNDER ACTIVE DEVELOPMENT and are not yet ready for production use. -# They will be used for the upcoming distributed version of chroma. They are not even ready -# for testing yet. Please do not use them unless you are working on the distributed version of chroma. - -apiVersion: v1 -kind: Namespace -metadata: - name: chroma - ---- - -apiVersion: v1 -kind: Service -metadata: - name: pulsar - namespace: chroma -spec: - ports: - - name: pulsar-port - port: 6650 - targetPort: 6650 - - name: admin-port - port: 8080 - targetPort: 8080 - selector: - app: pulsar - type: ClusterIP - ---- - -# TODO: Should be stateful set locally or managed via terraform into streamnative for cloud deployment -apiVersion: apps/v1 -kind: Deployment -metadata: - name: pulsar - namespace: chroma -spec: - replicas: 1 - selector: - matchLabels: - app: pulsar - template: - metadata: - labels: - app: pulsar - spec: - containers: - - name: pulsar - image: apachepulsar/pulsar - command: [ "/pulsar/bin/pulsar", "standalone" ] - env: - # This is needed by github actions. We force this to be lower everywehre for now. - # Since real deployments will configure/use pulsar this way. - - name: PULSAR_MEM - value: "-Xms128m -Xmx512m" - ports: - - containerPort: 6650 - - containerPort: 8080 - volumeMounts: - - name: pulsardata - mountPath: /pulsar/data - # readinessProbe: - # httpGet: - # path: /admin/v2/brokers/health - # port: 8080 - # initialDelaySeconds: 10 - # periodSeconds: 5 - # livenessProbe: - # httpGet: - # path: /admin/v2/brokers/health - # port: 8080 - # initialDelaySeconds: 20 - # periodSeconds: 10 - volumes: - - name: pulsardata - emptyDir: {} - ---- - -apiVersion: v1 -kind: Service -metadata: - name: server - namespace: chroma -spec: - ports: - - name: server - port: 8000 - targetPort: 8000 - selector: - app: server - type: LoadBalancer - ---- - -apiVersion: apps/v1 -kind: Deployment -metadata: - name: server - namespace: chroma -spec: - replicas: 1 - selector: - matchLabels: - app: server - template: - metadata: - labels: - app: server - spec: - containers: - - name: server - image: server - imagePullPolicy: IfNotPresent - ports: - - containerPort: 8000 - volumeMounts: - - name: chroma - mountPath: /test - env: - - name: IS_PERSISTENT - value: "TRUE" - - name: CHROMA_PRODUCER_IMPL - value: "chromadb.ingest.impl.pulsar.PulsarProducer" - - name: CHROMA_CONSUMER_IMPL - value: "chromadb.ingest.impl.pulsar.PulsarConsumer" - - name: CHROMA_SEGMENT_MANAGER_IMPL - value: "chromadb.segment.impl.manager.distributed.DistributedSegmentManager" - - name: PULSAR_BROKER_URL - value: "pulsar.chroma" - - name: PULSAR_BROKER_PORT - value: "6650" - - name: PULSAR_ADMIN_PORT - value: "8080" - - name: ALLOW_RESET - value: "TRUE" - - name: CHROMA_SYSDB_IMPL - value: "chromadb.db.impl.grpc.client.GrpcSysDB" - - name: CHROMA_SERVER_GRPC_PORT - value: "50051" - - name: CHROMA_COORDINATOR_HOST - value: "coordinator.chroma" - readinessProbe: - httpGet: - path: /api/v1/heartbeat - port: 8000 - initialDelaySeconds: 10 - periodSeconds: 5 - # livenessProbe: - # httpGet: - # path: /healthz - # port: 8000 - # initialDelaySeconds: 20 - # periodSeconds: 10 - # Ephemeral for now - volumes: - - name: chroma - emptyDir: {} - ---- - -# apiVersion: v1 -# kind: PersistentVolumeClaim -# metadata: -# name: index-data -# namespace: chroma -# spec: -# accessModes: -# - ReadWriteOnce -# resources: -# requests: -# storage: 1Gi - -apiVersion: apps/v1 -kind: Deployment -metadata: - name: coordinator - namespace: chroma -spec: - replicas: 1 - selector: - matchLabels: - app: coordinator - template: - metadata: - labels: - app: coordinator - spec: - containers: - - command: - - "chroma" - - "coordinator" - - "--pulsar-admin-url=http://pulsar.chroma:8080" - - "--pulsar-url=pulsar://pulsar.chroma:6650" - - "--notifier-provider=pulsar" - image: chroma-coordinator - imagePullPolicy: IfNotPresent - name: coordinator - ports: - - containerPort: 50051 - name: grpc - resources: - limits: - cpu: 100m - memory: 128Mi - ---- - -apiVersion: v1 -kind: Service -metadata: - name: coordinator - namespace: chroma -spec: - ports: - - name: grpc - port: 50051 - targetPort: grpc - selector: - app: coordinator - type: ClusterIP diff --git a/k8s/deployment/segment-server.yaml b/k8s/deployment/segment-server.yaml deleted file mode 100644 index 33af91d1314..00000000000 --- a/k8s/deployment/segment-server.yaml +++ /dev/null @@ -1,87 +0,0 @@ -apiVersion: v1 -kind: Service -metadata: - name: segment-server - namespace: chroma -spec: - ports: - - name: segment-server-port - port: 50051 - targetPort: 50051 - selector: - app: segment-server - type: ClusterIP - ---- - -apiVersion: apps/v1 -kind: Deployment -metadata: - name: segment-server - namespace: chroma -spec: - replicas: 1 - selector: - matchLabels: - app: segment-server - template: - metadata: - labels: - app: segment-server - member-type: worker - spec: - containers: - - name: segment-server - image: worker - imagePullPolicy: IfNotPresent - command: ["cargo", "run"] - ports: - - containerPort: 50051 - volumeMounts: - - name: chroma - mountPath: /index_data - env: - - name: CHROMA_WORKER__PULSAR_URL - value: pulsar://pulsar.chroma:6650 - - name: CHROMA_WORKER__PULSAR_NAMESPACE - value: default - - name: CHROMA_WORKER__PULSAR_TENANT - value: default - - name: CHROMA_WORKER__MY_IP - valueFrom: - fieldRef: - fieldPath: status.podIP - # livenessProbe: - # grpc: - # port: 50051 - # initialDelaySeconds: 10 - volumes: - - name: chroma - emptyDir: {} - ---- - -apiVersion: rbac.authorization.k8s.io/v1 -kind: Role -metadata: - namespace: chroma - name: pod-watcher -rules: -- apiGroups: [""] - resources: ["pods"] - verbs: ["get", "list", "watch"] - ---- -apiVersion: rbac.authorization.k8s.io/v1 -kind: RoleBinding -metadata: - name: pod-watcher-binding - namespace: chroma -subjects: -- kind: ServiceAccount - name: default - namespace: chroma -roleRef: - kind: Role - name: pod-watcher - apiGroup: rbac.authorization.k8s.io diff --git a/k8s/distributed-chroma/.helmignore b/k8s/distributed-chroma/.helmignore new file mode 100644 index 00000000000..0e8a0eb36f4 --- /dev/null +++ b/k8s/distributed-chroma/.helmignore @@ -0,0 +1,23 @@ +# Patterns to ignore when building packages. +# This supports shell glob matching, relative path matching, and +# negation (prefixed with !). Only one pattern per line. +.DS_Store +# Common VCS dirs +.git/ +.gitignore +.bzr/ +.bzrignore +.hg/ +.hgignore +.svn/ +# Common backup files +*.swp +*.bak +*.tmp +*.orig +*~ +# Various IDEs +.project +.idea/ +*.tmproj +.vscode/ diff --git a/k8s/distributed-chroma/Chart.yaml b/k8s/distributed-chroma/Chart.yaml new file mode 100644 index 00000000000..f0a51511dd2 --- /dev/null +++ b/k8s/distributed-chroma/Chart.yaml @@ -0,0 +1,30 @@ +# Copyright 2024 Chroma Inc +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +apiVersion: v2 +name: distributed-chroma +description: A helm chart for distributed Chroma +type: application +version: 0.1.2 +appVersion: "0.4.23" +keywords: + - chroma + - vector + - database + - retrieval + - llm + - rag +home: "https://www.trychroma.com/" +sources: + - "https://github.com/chroma-core/chroma" diff --git a/k8s/crd/memberlist_crd.yaml b/k8s/distributed-chroma/crds/memberlist_crd.yaml similarity index 92% rename from k8s/crd/memberlist_crd.yaml rename to k8s/distributed-chroma/crds/memberlist_crd.yaml index 9d31706aad2..fb593c7e2a7 100644 --- a/k8s/crd/memberlist_crd.yaml +++ b/k8s/distributed-chroma/crds/memberlist_crd.yaml @@ -2,6 +2,7 @@ # They will be used for the upcoming distributed version of chroma. They are not even ready # for testing yet. Please do not use them unless you are working on the distributed version of chroma. +# Note from ben: Before you modify this please read https://hackmd.io/@carvel/rJKraqlDD apiVersion: apiextensions.k8s.io/v1 kind: CustomResourceDefinition metadata: diff --git a/k8s/distributed-chroma/templates/compaction-service-memberlist-cr.yaml b/k8s/distributed-chroma/templates/compaction-service-memberlist-cr.yaml new file mode 100644 index 00000000000..b68981c05ee --- /dev/null +++ b/k8s/distributed-chroma/templates/compaction-service-memberlist-cr.yaml @@ -0,0 +1,79 @@ +# These kubernetes manifests are UNDER ACTIVE DEVELOPMENT and are not yet ready for production use. +# They will be used for the upcoming distributed version of chroma. They are not even ready +# for testing yet. Please do not use them unless you are working on the distributed version of chroma. + +apiVersion: chroma.cluster/v1 +kind: MemberList +metadata: + name: compaction-service-memberlist + namespace: {{ .Values.namespace}} +spec: + members: + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: compaction-service-memberlist-readerwriter +rules: +- apiGroups: + - chroma.cluster + resources: + - memberlists + verbs: + - get + - list + - watch + # TODO: FIX THIS LEAKY PERMISSION + - create + - update + - patch + - delete + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: sysdb-compaction-service-memberlist-binding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: compaction-service-memberlist-readerwriter +subjects: +- kind: ServiceAccount + name: sysdb-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + # Awkward name, but this lets the compaction-service-serviceaccount read + # the compaction-service-memberlist. + name: compaction-service-compaction-service-memberlist-binding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: compaction-service-memberlist-readerwriter +subjects: +- kind: ServiceAccount + name: compaction-service-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: compaction-service-memberlist-readerwriter-binding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: compaction-service-memberlist-readerwriter +subjects: +- kind: ServiceAccount + name: default + namespace: {{ .Values.namespace }} diff --git a/k8s/distributed-chroma/templates/compaction-service.yaml b/k8s/distributed-chroma/templates/compaction-service.yaml new file mode 100644 index 00000000000..5fd39e382e8 --- /dev/null +++ b/k8s/distributed-chroma/templates/compaction-service.yaml @@ -0,0 +1,64 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: compaction-service + namespace: {{ .Values.namespace }} +spec: + replicas: 1 + selector: + matchLabels: + app: compaction-service + template: + metadata: + labels: + app: compaction-service + member-type: compaction-service + spec: + serviceAccountName: compaction-service-serviceaccount + containers: + - name: compaction-service + image: "{{ .Values.compactionService.image.repository }}:{{ .Values.compactionService.image.tag }}" + imagePullPolicy: IfNotPresent + ports: + - containerPort: 50051 + env: + - name: CHROMA_compaction-service__MY_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: "kubernetes.io/hostname" + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + member-type: compaction-service + volumes: + - name: chroma + emptyDir: {} + +--- + +apiVersion: v1 +kind: ServiceAccount +metadata: + name: compaction-service-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: compaction-service-serviceaccount-rolebinding + namespace: {{ .Values.namespace }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: pod-watcher +subjects: +- kind: ServiceAccount + name: compaction-service-serviceaccount + namespace: {{ .Values.namespace }} + +--- diff --git a/k8s/distributed-chroma/templates/frontend-service.yaml b/k8s/distributed-chroma/templates/frontend-service.yaml new file mode 100644 index 00000000000..067fd033dea --- /dev/null +++ b/k8s/distributed-chroma/templates/frontend-service.yaml @@ -0,0 +1,81 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: frontend-service + namespace: {{ .Values.namespace }} +spec: + replicas: 2 + selector: + matchLabels: + app: frontend-service + template: + metadata: + labels: + app: frontend-service + spec: + containers: + - name: frontend-service + image: "{{ .Values.frontendService.image.repository }}:{{ .Values.frontendService.image.tag }}" + imagePullPolicy: IfNotPresent + ports: + - containerPort: 8000 + volumeMounts: + - name: chroma + mountPath: /test + env: + - name: IS_PERSISTENT + {{ .Values.frontendService.isPersistent }} + - name: CHROMA_PRODUCER_IMPL + {{ .Values.frontendService.producerImpl }} + - name: CHROMA_CONSUMER_IMPL + {{ .Values.frontendService.consumerImpl }} + - name: CHROMA_SEGMENT_MANAGER_IMPL + {{ .Values.frontendService.segmentManagerImpl }} + - name: PULSAR_BROKER_URL + {{ .Values.frontendService.pulsarBrokerUrl }} + - name: PULSAR_BROKER_PORT + {{ .Values.frontendService.pulsarBrokerPort }} + - name: PULSAR_ADMIN_PORT + {{ .Values.frontendService.pulsarAdminPort }} + - name: ALLOW_RESET + {{ .Values.frontendService.allowReset }} + - name: CHROMA_SYSDB_IMPL + {{ .Values.frontendService.sysdbImpl }} + - name: CHROMA_SERVER_GRPC_PORT + {{ .Values.frontendService.serverGrpcPort }} + - name: CHROMA_COORDINATOR_HOST + {{ .Values.frontendService.coordinatorHost }} + - name: CHROMA_SERVER_AUTH_PROVIDER + {{ .Values.frontendService.authProvider }} + - name: CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER + {{ .Values.frontendService.authCredentialsProvider }} + - name: CHROMA_SERVER_AUTHZ_PROVIDER + {{ .Values.frontendService.authzProvider }} + - name: CHROMA_SERVER_AUTHZ_CONFIG_PROVIDER + {{ .Values.frontendService.authzConfigProvider }} + - name: CHROMA_MEMBERLIST_PROVIDER_IMPL + {{ .Values.frontendService.memberlistProviderImpl }} + - name: CHROMA_LOGSERVICE_HOST + {{ .Values.frontendService.logServiceHost }} + - name: CHROMA_LOGSERVICE_PORT + {{ .Values.frontendService.logServicePort }} +{{ .Values.frontendService.otherEnvConfig | nindent 12 }} + volumes: + - name: chroma + emptyDir: {} + +--- + +apiVersion: v1 +kind: Service +metadata: + name: frontend-service + namespace: {{ .Values.namespace }} +spec: + ports: + - name: server-port + port: 8000 + targetPort: 8000 + selector: + app: frontend-service + type: ClusterIP diff --git a/k8s/distributed-chroma/templates/logservice.yaml b/k8s/distributed-chroma/templates/logservice.yaml new file mode 100644 index 00000000000..88e4c4aedad --- /dev/null +++ b/k8s/distributed-chroma/templates/logservice.yaml @@ -0,0 +1,58 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: logservice + namespace: {{ .Values.namespace }} +spec: + replicas: 1 + selector: + matchLabels: + app: logservice + template: + metadata: + labels: + app: logservice + spec: + serviceAccountName: logservice-serviceaccount + containers: + - command: + - "/bin/sh" + - "-c" + # This has to be one line to be passed into the `exec` env correctly. I truly could not tell you why. + - logservice logservice {{ range $k, $v := .Values.logService.flags }} --{{ $k }}={{ $v }} {{ end }} + env: + {{ range .Values.logService.env }} + - name: {{ .name }} + # TODO properly use flow control here to check which type of value we need. +{{ .value | nindent 14 }} + {{ end }} + image: "{{ .Values.logService.image.repository }}:{{ .Values.logService.image.tag }}" + imagePullPolicy: IfNotPresent + name: logservice + ports: + - containerPort: 50051 + name: grpc +--- +apiVersion: v1 +kind: Service +metadata: + name: logservice + namespace: {{ .Values.namespace }} +spec: + ports: + - name: grpc + port: 50051 + targetPort: grpc + selector: + app: logservice + type: ClusterIP + +--- + +apiVersion: v1 +kind: ServiceAccount +metadata: + name: logservice-serviceaccount + namespace: {{ .Values.namespace }} + +--- diff --git a/k8s/distributed-chroma/templates/namespace.yaml b/k8s/distributed-chroma/templates/namespace.yaml new file mode 100644 index 00000000000..48685640e18 --- /dev/null +++ b/k8s/distributed-chroma/templates/namespace.yaml @@ -0,0 +1,8 @@ +--- + +apiVersion: v1 +kind: Namespace +metadata: + name: {{ .Values.namespace }} + +--- \ No newline at end of file diff --git a/k8s/distributed-chroma/templates/pod-watcher-role.yaml b/k8s/distributed-chroma/templates/pod-watcher-role.yaml new file mode 100644 index 00000000000..eb8ff467961 --- /dev/null +++ b/k8s/distributed-chroma/templates/pod-watcher-role.yaml @@ -0,0 +1,13 @@ +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + namespace: {{ .Values.namespace }} + name: pod-watcher +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "watch"] + +--- \ No newline at end of file diff --git a/k8s/distributed-chroma/templates/pulsar.yaml b/k8s/distributed-chroma/templates/pulsar.yaml new file mode 100644 index 00000000000..8d487650335 --- /dev/null +++ b/k8s/distributed-chroma/templates/pulsar.yaml @@ -0,0 +1,51 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: pulsar + namespace: {{ .Values.namespace }} +spec: + replicas: 1 + selector: + matchLabels: + app: pulsar + template: + metadata: + labels: + app: pulsar + spec: + containers: + - name: pulsar + image: apachepulsar/pulsar + command: [ "/pulsar/bin/pulsar", "standalone" ] + ports: + - containerPort: 6650 + - containerPort: 8080 + volumeMounts: + - name: pulsardata + mountPath: /pulsar/data + readinessProbe: + httpGet: + path: /admin/v2/brokers/health + port: 8080 + initialDelaySeconds: 10 + periodSeconds: 5 + volumes: + - name: pulsardata + emptyDir: {} +--- +apiVersion: v1 +kind: Service +metadata: + name: pulsar + namespace: {{ .Values.namespace }} +spec: + ports: + - name: pulsar-port + port: 6650 + targetPort: 6650 + - name: admin-port + port: 8080 + targetPort: 8080 + selector: + app: pulsar + type: ClusterIP \ No newline at end of file diff --git a/k8s/distributed-chroma/templates/query-service.yaml b/k8s/distributed-chroma/templates/query-service.yaml new file mode 100644 index 00000000000..674ec8d44bb --- /dev/null +++ b/k8s/distributed-chroma/templates/query-service.yaml @@ -0,0 +1,87 @@ +--- + +apiVersion: v1 +kind: Service +metadata: + name: query-service + namespace: {{ .Values.namespace }} +spec: + ports: + - name: query-service-server-port + port: 50051 + targetPort: 50051 + selector: + app: query-service-server + type: ClusterIP + +--- + +apiVersion: apps/v1 +kind: Deployment +metadata: + name: query-service + namespace: {{ .Values.namespace }} +spec: + replicas: 2 + selector: + matchLabels: + app: query-service + template: + metadata: + labels: + app: query-service + member-type: query-service + spec: + serviceAccountName: query-service-serviceaccount + containers: + - name: query-service + image: "{{ .Values.queryService.image.repository }}:{{ .Values.queryService.image.tag }}" + imagePullPolicy: IfNotPresent + ports: + - containerPort: 50051 + volumeMounts: + - name: chroma + mountPath: /index_data + env: + - name: CHROMA_query-service__PULSAR_URL + value: pulsar://pulsar.chroma:6650 + - name: CHROMA_query-service__MY_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: "kubernetes.io/hostname" + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + member-type: query-service + volumes: + - name: chroma + emptyDir: {} + +--- + +apiVersion: v1 +kind: ServiceAccount +metadata: + name: query-service-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: query-service-serviceaccount-rolebinding + namespace: {{ .Values.namespace }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: pod-watcher +subjects: +- kind: ServiceAccount + name: query-service-serviceaccount + namespace: {{ .Values.namespace }} + +--- \ No newline at end of file diff --git a/k8s/distributed-chroma/templates/sysdb-migration.yaml b/k8s/distributed-chroma/templates/sysdb-migration.yaml new file mode 100644 index 00000000000..dc62d157603 --- /dev/null +++ b/k8s/distributed-chroma/templates/sysdb-migration.yaml @@ -0,0 +1,28 @@ +apiVersion: batch/v1 +kind: Job +metadata: + name: sysdb-migration + namespace: {{ .Values.namespace }} +spec: + template: + metadata: + labels: + app: sysdb-migration + spec: + restartPolicy: OnFailure + containers: + - command: + - "/bin/sh" + - "-c" + - "atlas migrate apply --url postgres://{{ .Values.sysdbMigration.username }}:{{ .Values.sysdbMigration.password }}@{{ .Values.sysdbMigration.netloc }}:{{ .Values.sysdbMigration.port }}/{{ .Values.sysdbMigration.dbName }}?sslmode={{ .Values.sysdbMigration.sslmode }}" + image: "{{ .Values.sysdbMigration.image.repository }}:{{ .Values.sysdbMigration.image.tag }}" + imagePullPolicy: IfNotPresent + name: migration + env: + {{ range .Values.sysdb.env }} + - name: {{ .name }} + # TODO properly use flow control here to check which type of value we need. +{{ .value | nindent 14 }} + {{ end }} + +--- diff --git a/k8s/distributed-chroma/templates/sysdb-service.yaml b/k8s/distributed-chroma/templates/sysdb-service.yaml new file mode 100644 index 00000000000..296d391c485 --- /dev/null +++ b/k8s/distributed-chroma/templates/sysdb-service.yaml @@ -0,0 +1,76 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: sysdb + namespace: {{ .Values.namespace }} +spec: + replicas: {{ .Values.sysdb.replicaCount }} + selector: + matchLabels: + app: sysdb + template: + metadata: + labels: + app: sysdb + spec: + serviceAccountName: sysdb-serviceaccount + containers: + - command: + - "/bin/sh" + - "-c" + # This has to be one line to be passed into the `exec` env correctly. I truly could not tell you why. + - coordinator coordinator {{ range $k, $v := .Values.sysdb.flags }} --{{ $k }}={{ $v }} {{ end }} + env: + {{ range .Values.sysdb.env }} + - name: {{ .name }} + # TODO properly use flow control here to check which type of value we need. +{{ .value | nindent 14 }} + {{ end }} + image: "{{ .Values.sysdb.image.repository }}:{{ .Values.sysdb.image.tag }}" + imagePullPolicy: IfNotPresent + name: sysdb + ports: + - containerPort: 50051 + name: grpc + +--- + +apiVersion: v1 +kind: Service +metadata: + name: sysdb + namespace: {{ .Values.namespace }} +spec: + ports: + - name: grpc + port: 50051 + targetPort: grpc + selector: + app: sysdb + type: ClusterIP + +--- + +apiVersion: v1 +kind: ServiceAccount +metadata: + name: sysdb-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: sysdb-serviceaccount-rolebinding + namespace: {{ .Values.namespace }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: pod-watcher +subjects: +- kind: ServiceAccount + name: sysdb-serviceaccount + namespace: {{ .Values.namespace }} + +--- diff --git a/k8s/distributed-chroma/templates/worker-memberlist-cr.yaml b/k8s/distributed-chroma/templates/worker-memberlist-cr.yaml new file mode 100644 index 00000000000..2261057ec5d --- /dev/null +++ b/k8s/distributed-chroma/templates/worker-memberlist-cr.yaml @@ -0,0 +1,79 @@ +# These kubernetes manifests are UNDER ACTIVE DEVELOPMENT and are not yet ready for production use. +# They will be used for the upcoming distributed version of chroma. They are not even ready +# for testing yet. Please do not use them unless you are working on the distributed version of chroma. + +apiVersion: chroma.cluster/v1 +kind: MemberList +metadata: + name: query-service-memberlist + namespace: {{ .Values.namespace}} +spec: + members: + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: query-service-memberlist-readerwriter +rules: +- apiGroups: + - chroma.cluster + resources: + - memberlists + verbs: + - get + - list + - watch + # TODO: FIX THIS LEAKY PERMISSION + - create + - update + - patch + - delete + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: sysdb-query-service-memberlist-binding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: query-service-memberlist-readerwriter +subjects: +- kind: ServiceAccount + name: sysdb-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + # Awkward name, but this lets the query-service-serviceaccount read + # the query-service-memberlist. + name: query-service-query-service-memberlist-binding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: query-service-memberlist-readerwriter +subjects: +- kind: ServiceAccount + name: query-service-serviceaccount + namespace: {{ .Values.namespace }} + +--- + +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: query-service-memberlist-readerwriter-binding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: query-service-memberlist-readerwriter +subjects: +- kind: ServiceAccount + name: default + namespace: {{ .Values.namespace }} diff --git a/k8s/distributed-chroma/values.yaml b/k8s/distributed-chroma/values.yaml new file mode 100644 index 00000000000..a7be3998c0d --- /dev/null +++ b/k8s/distributed-chroma/values.yaml @@ -0,0 +1,72 @@ +# Default values for distributed-chroma. +# Strongly prefer single quotes. + +namespace: 'chroma' + +frontendService: + image: + repository: 'local' + tag: 'frontend-service' + + # Sometimes users (and us) want to pass values directly as flags. Sometimes, these are + # populated from secrets or configMaps. So we let consumers fill these directly. + # TODO we could maybe have mutually exclusive pieces of config, e.g. isPersistentValue + # and isPersistentFromConfigMap or something. + isPersistent: 'value: "TRUE"' + producerImpl: 'value: "chromadb.logservice.logservice.LogService"' + consumerImpl: 'value: "chromadb.logservice.logservice.LogService"' + segmentManagerImpl: 'value: "chromadb.segment.impl.manager.distributed.DistributedSegmentManager"' + pulsarBrokerUrl: 'value: "pulsar.chroma"' + pulsarBrokerPort: 'value: "6650"' + pulsarAdminPort: 'value: "8080"' + allowReset: 'value: "TRUE"' + sysdbImpl: 'value: "chromadb.db.impl.grpc.client.GrpcSysDB"' + serverGrpcPort: 'value: "50051"' + coordinatorHost: 'value: "sysdb.chroma"' + authProvider: 'value: ""' + authCredentialsProvider: 'value: ""' + authzProvider: 'value: ""' + authzConfigProvider: 'value: ""' + memberlistProviderImpl: 'value: "chromadb.segment.impl.distributed.segment_directory.MockMemberlistProvider"' + logServiceHost: 'value: "logservice.chroma"' + logServicePort: 'value: "50051"' + otherEnvConfig: '' + +sysdb: + image: + repository: 'local' + tag: 'sysdb' + replicaCount: 1 + env: + flags: + pulsar-admin-url: "http://pulsar.chroma:8080" + pulsar-url: "pulsar://pulsar.chroma:6650" + notifier-provider: "pulsar" + +logService: + image: + repository: 'local' + tag: 'sysdb' + env: + flags: + +queryService: + image: + repository: 'local' + tag: 'query-service' + +compactionService: + image: + repository: 'local' + tag: 'compaction-service' + +sysdbMigration: + image: + repository: 'local' + tag: 'sysdb-migration' + username: chroma + password: chroma + netloc: postgres + port: 5432 + dbName: chroma + sslmode: disable diff --git a/k8s/test/README.md b/k8s/test/README.md new file mode 100644 index 00000000000..83ffcc949a5 --- /dev/null +++ b/k8s/test/README.md @@ -0,0 +1 @@ +This directory contains kubernetes manifests to be applied on top of our production manifests to make testing and debugging easier. For example, service endpoints to expose internal services. \ No newline at end of file diff --git a/k8s/test/jaeger-service.yaml b/k8s/test/jaeger-service.yaml new file mode 100644 index 00000000000..cfb95ed3036 --- /dev/null +++ b/k8s/test/jaeger-service.yaml @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: Service +metadata: + name: jaeger-lb + namespace: chroma +spec: + ports: + - name: http + port: 16686 + targetPort: 16686 + selector: + app: jaeger + type: LoadBalancer \ No newline at end of file diff --git a/k8s/test/jaeger.yaml b/k8s/test/jaeger.yaml new file mode 100644 index 00000000000..02dff6a5022 --- /dev/null +++ b/k8s/test/jaeger.yaml @@ -0,0 +1,48 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: jaeger + namespace: {{ .Values.namespace }} +spec: + replicas: 1 + selector: + matchLabels: + app: jaeger + template: + metadata: + labels: + app: jaeger + spec: + containers: + - name: jaeger + image: jaegertracing/all-in-one:1.35 + env: + - name: COLLECTOR_OTLP_ENABLED + value: "true" + ports: + - containerPort: 16686 + name: ui-port + - containerPort: 4317 + name: grpc-port + - containerPort: 4318 + name: http-port + + +--- +apiVersion: v1 +kind: Service +metadata: + name: jaeger + namespace: {{ .Values.namespace }} +spec: + type: ClusterIP + ports: + - port: 16686 + name: ui-port + - port: 4317 + name: grpc-port + - port: 4318 + name: http-port + selector: + app: jaeger + diff --git a/k8s/test/coordinator_service.yaml b/k8s/test/logservice-service.yaml similarity index 79% rename from k8s/test/coordinator_service.yaml rename to k8s/test/logservice-service.yaml index 37334b12187..2d3a4a8566a 100644 --- a/k8s/test/coordinator_service.yaml +++ b/k8s/test/logservice-service.yaml @@ -1,7 +1,7 @@ apiVersion: v1 kind: Service metadata: - name: coordinator-lb + name: logservice-lb namespace: chroma spec: ports: @@ -9,5 +9,5 @@ spec: port: 50051 targetPort: 50051 selector: - app: coordinator + app: logservice type: LoadBalancer diff --git a/k8s/test/minio.yaml b/k8s/test/minio.yaml index 148c5170fd8..d535e896234 100644 --- a/k8s/test/minio.yaml +++ b/k8s/test/minio.yaml @@ -18,25 +18,25 @@ spec: - name: minio emptyDir: {} containers: - - name: minio - image: minio/minio:latest - args: - - server - - /storage - env: - - name: MINIO_ACCESS_KEY - value: "minio" - - name: MINIO_SECRET_KEY - value: "minio123" - ports: - - containerPort: 9000 - hostPort: 9000 - volumeMounts: - name: minio - mountPath: /storage + image: minio/minio:latest + args: + - server + - /storage + env: + - name: MINIO_ACCESS_KEY + value: "minio" + - name: MINIO_SECRET_KEY + value: "minio123" + ports: + - containerPort: 9000 + hostPort: 9000 + name: http + volumeMounts: + - name: minio + mountPath: /storage --- - apiVersion: v1 kind: Service metadata: diff --git a/k8s/test/postgres.yaml b/k8s/test/postgres.yaml new file mode 100644 index 00000000000..33d3417a340 --- /dev/null +++ b/k8s/test/postgres.yaml @@ -0,0 +1,43 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: postgres + namespace: chroma +spec: + replicas: 1 + selector: + matchLabels: + app: postgres + template: + metadata: + labels: + app: postgres + spec: + containers: + - name: postgres + image: postgres:14.1-alpine + env: + - name: POSTGRES_DB + value: chroma + - name: POSTGRES_USER + value: chroma + - name: POSTGRES_PASSWORD + value: chroma + ports: + - containerPort: 5432 + +--- + +apiVersion: v1 +kind: Service +metadata: + name: postgres + namespace: chroma +spec: + ports: + - name: postgres-port + port: 5432 + targetPort: 5432 + selector: + app: postgres + type: ClusterIP diff --git a/k8s/test/pulsar_service.yaml b/k8s/test/pulsar-service.yaml similarity index 100% rename from k8s/test/pulsar_service.yaml rename to k8s/test/pulsar-service.yaml diff --git a/k8s/test/segment_server_service.yml b/k8s/test/query-service-service.yaml similarity index 64% rename from k8s/test/segment_server_service.yml rename to k8s/test/query-service-service.yaml index 7463333deef..6cfdd71677b 100644 --- a/k8s/test/segment_server_service.yml +++ b/k8s/test/query-service-service.yaml @@ -1,13 +1,13 @@ apiVersion: v1 kind: Service metadata: - name: segment-server-lb + name: query-service-lb namespace: chroma spec: ports: - - name: segment-server-port + - name: query-service-port port: 50052 targetPort: 50051 selector: - app: segment-server + app: query-service type: LoadBalancer diff --git a/k8s/test/sysdb-service.yaml b/k8s/test/sysdb-service.yaml new file mode 100644 index 00000000000..dda139fcb7d --- /dev/null +++ b/k8s/test/sysdb-service.yaml @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: Service +metadata: + name: sysdb-lb + namespace: chroma +spec: + ports: + - name: grpc + port: 50051 + targetPort: 50051 + selector: + app: sysdb + type: LoadBalancer diff --git a/k8s/test/test_memberlist_cr.yaml b/k8s/test/test-memberlist-cr.yaml similarity index 100% rename from k8s/test/test_memberlist_cr.yaml rename to k8s/test/test-memberlist-cr.yaml diff --git a/pull_request_template.md b/pull_request_template.md index b7fbdce6fc6..3e1091ff328 100644 --- a/pull_request_template.md +++ b/pull_request_template.md @@ -9,7 +9,7 @@ ## Test plan *How are these changes tested?* -- [ ] Tests pass locally with `pytest` for python, `yarn test` for js +- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* diff --git a/pyproject.toml b/pyproject.toml index 01aa0d8663b..8e5c29527e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,6 @@ dependencies = [ 'numpy >= 1.22.5', 'posthog >= 2.4.0', 'typing_extensions >= 4.5.0', - 'pulsar-client>=3.1.0', 'onnxruntime >= 1.14.1', 'opentelemetry-api>=1.2.0', 'opentelemetry-exporter-otlp-proto-grpc>=1.2.0', @@ -43,6 +42,7 @@ dependencies = [ 'tenacity>=8.2.3', 'PyYAML>=6.0.0', 'mmh3>=4.0.1', + 'orjson>=3.9.12', ] [tool.black] diff --git a/requirements.txt b/requirements.txt index 3e99734fd3a..02e7c2a62bb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ -bcrypt==4.0.1 -chroma-hnswlib==0.7.3 +bcrypt>=4.0.1 +chroma-hnswlib>=0.7.3 fastapi>=0.95.2 graphlib_backport==1.0.3; python_version < '3.9' grpcio>=1.58.0 @@ -12,16 +12,16 @@ opentelemetry-api>=1.2.0 opentelemetry-exporter-otlp-proto-grpc>=1.2.0 opentelemetry-instrumentation-fastapi>=0.41b0 opentelemetry-sdk>=1.2.0 -overrides==7.3.1 -posthog==2.4.0 -pulsar-client==3.1.0 +orjson>=3.9.12 +overrides>=7.3.1 +posthog>=2.4.0 pydantic>=1.9 -pypika==0.48.9 +pypika>=0.48.9 PyYAML>=6.0.0 -requests==2.28.1 +requests>=2.28.1 tenacity>=8.2.3 -tokenizers==0.13.2 +tokenizers>=0.13.2 tqdm>=4.65.0 typer>=0.9.0 typing_extensions>=4.5.0 -uvicorn[standard]==0.18.3 +uvicorn[standard]>=0.18.3 diff --git a/requirements_dev.txt b/requirements_dev.txt index 4dce86e2efe..4df73fb8c56 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -2,8 +2,8 @@ black==23.3.0 # match what's in pyproject.toml build grpcio-tools httpx -hypothesis -hypothesis[numpy] +hypothesis<=6.98.9 # > Than this version has API changes we don't currently support +hypothesis[numpy]<=6.98.9 mypy-protobuf pre-commit pytest diff --git a/rust/worker/.dockerignore b/rust/worker/.dockerignore new file mode 100644 index 00000000000..2f0f51534e8 --- /dev/null +++ b/rust/worker/.dockerignore @@ -0,0 +1,3 @@ +.dockerignore +.gitignore +Dockerfile \ No newline at end of file diff --git a/rust/worker/Cargo.toml b/rust/worker/Cargo.toml index 25a3b2d099e..8f5f7c63e8e 100644 --- a/rust/worker/Cargo.toml +++ b/rust/worker/Cargo.toml @@ -4,8 +4,12 @@ version = "0.1.0" edition = "2021" [[bin]] -name = "worker" -path = "src/bin/worker.rs" +name = "query_service" +path = "src/bin/query_service.rs" + +[[bin]] +name = "compaction_service" +path = "src/bin/compaction_service.rs" [dependencies] tonic = "0.10" @@ -25,7 +29,6 @@ num_cpus = "1.16.0" pulsar = "6.1.0" murmur3 = "0.5.2" thiserror = "1.0.50" -num-bigint = "0.4.4" tempfile = "3.8.1" schemars = "0.8.16" kube = { version = "0.87.1", features = ["runtime", "derive"] } @@ -35,6 +38,14 @@ parking_lot = "0.12.1" aws-sdk-s3 = "1.5.0" aws-smithy-types = "1.1.0" aws-config = { version = "1.1.2", features = ["behavior-version-latest"] } +arrow = "50.0.0" +roaring = "0.10.3" +tantivy = "0.21.1" + +[dev-dependencies] +proptest = "1.4.0" +proptest-state-machine = "0.1.0" +"rand" = "0.8.5" [build-dependencies] tonic-build = "0.10" diff --git a/rust/worker/Dockerfile b/rust/worker/Dockerfile index 2e3802787e1..91fea920990 100644 --- a/rust/worker/Dockerfile +++ b/rust/worker/Dockerfile @@ -2,20 +2,60 @@ FROM rust:1.74.1 as builder ARG CHROMA_KUBERNETES_INTEGRATION=0 ENV CHROMA_KUBERNETES_INTEGRATION $CHROMA_KUBERNETES_INTEGRATION +ARG RELEASE_MODE= + WORKDIR / RUN git clone https://github.com/chroma-core/hnswlib.git +# Cache dependencies by building them without our code first. +# https://dev.to/rogertorres/first-steps-with-docker-rust-30oi +# https://www.reddit.com/r/rust/comments/126xeyx/exploring_the_problem_of_faster_cargo_docker/ WORKDIR /chroma/ -COPY . . - +COPY Cargo.toml Cargo.toml +COPY Cargo.lock Cargo.lock +COPY idl/ idl/ +COPY /rust/worker/Cargo.toml rust/worker/Cargo.toml ENV PROTOC_ZIP=protoc-25.1-linux-x86_64.zip RUN curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v25.1/$PROTOC_ZIP \ && unzip -o $PROTOC_ZIP -d /usr/local bin/protoc \ && unzip -o $PROTOC_ZIP -d /usr/local 'include/*' \ && rm -f $PROTOC_ZIP -RUN cargo build +FROM builder as query_service_builder +# We need to replace the query node's real main() with a dummy at the expected location. +RUN mkdir -p rust/worker/src/bin/ && echo "fn main() {}" > rust/worker/src/bin/query_service.rs +RUN if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin query_service --release; else cargo build --bin query_service; fi +RUN rm -rf rust/ + +COPY rust/ rust/ +RUN touch rust/worker/src/bin/query_service.rs +RUN if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin query_service --release; else cargo build --bin query_service; fi +RUN if [ "$RELEASE_MODE" = "1" ]; then mv target/release/query_service .; else mv target/debug/query_service .; fi + +FROM debian:bookworm-slim as query_service + +COPY --from=query_service_builder /chroma/query_service . +COPY --from=query_service_builder /chroma/rust/worker/chroma_config.yaml . +RUN apt-get update && apt-get install -y libssl-dev + +ENTRYPOINT [ "./query_service" ] + +FROM builder as compaction_service_builder + +# We need to replace the compaction node's real main() with a dummy at the expected location. +RUN mkdir -p rust/worker/src/bin/ && echo "fn main() {}" > rust/worker/src/bin/compaction_service.rs +RUN if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin compaction_service --release; else cargo build --bin compaction_service; fi +RUN rm -rf rust/ + +COPY rust/ rust/ +RUN touch rust/worker/src/bin/compaction_service.rs +RUN if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin compaction_service --release; else cargo build --bin compaction_service; fi +RUN if [ "$RELEASE_MODE" = "1" ]; then mv target/release/compaction_service .; else mv target/debug/compaction_service .; fi + +FROM debian:bookworm-slim as compaction_service -WORKDIR /chroma/rust/worker +COPY --from=compaction_service_builder /chroma/compaction_service . +COPY --from=compaction_service_builder /chroma/rust/worker/chroma_config.yaml . +RUN apt-get update && apt-get install -y libssl-dev -CMD ["cargo", "run"] +ENTRYPOINT [ "./compaction_service" ] diff --git a/rust/worker/README b/rust/worker/README deleted file mode 100644 index e09a7db4f4c..00000000000 --- a/rust/worker/README +++ /dev/null @@ -1,7 +0,0 @@ -# Readme - - - - -### Rust version -Use rust 1.74.0 or greater. diff --git a/rust/worker/README.md b/rust/worker/README.md new file mode 100644 index 00000000000..421ae6543b1 --- /dev/null +++ b/rust/worker/README.md @@ -0,0 +1,15 @@ +# Readme + +This folder houses the Rust code for the query and compactor nodes. It is a standard rust crate managed using cargo. + +### Testing + +`cargo test` + +### Building + +`cargo build` + +### Rust version + +Use rust 1.74.0 or greater. diff --git a/rust/worker/build.rs b/rust/worker/build.rs index 25235b5c6b0..c98a96de683 100644 --- a/rust/worker/build.rs +++ b/rust/worker/build.rs @@ -1,12 +1,15 @@ fn main() -> Result<(), Box> { // Compile the protobuf files in the chromadb proto directory. - tonic_build::configure().compile( - &[ - "../../idl/chromadb/proto/chroma.proto", - "../../idl/chromadb/proto/coordinator.proto", - ], - &["../../idl/"], - )?; + tonic_build::configure() + .emit_rerun_if_changed(true) + .compile( + &[ + "../../idl/chromadb/proto/chroma.proto", + "../../idl/chromadb/proto/coordinator.proto", + "../../idl/chromadb/proto/logservice.proto", + ], + &["../../idl/"], + )?; // Compile the hnswlib bindings. cc::Build::new() @@ -17,6 +20,7 @@ fn main() -> Result<(), Box> { .flag("-DHAVE_CXX0X") .flag("-fpic") .flag("-ftree-vectorize") + .flag("-w") .compile("bindings"); // Set a compile flag based on an environment variable that tells us if we should diff --git a/rust/worker/chroma_config.yaml b/rust/worker/chroma_config.yaml index 32e5165924d..40e6e1799c3 100644 --- a/rust/worker/chroma_config.yaml +++ b/rust/worker/chroma_config.yaml @@ -1,31 +1,62 @@ -# Default configuration for Chroma worker +# Default configuration for query and compaction service # In the long term, every service should have an entry in this file # and this can become the global configuration file for Chroma # for now we nest it in the worker directory -worker: +query_service: my_ip: "10.244.0.9" my_port: 50051 - num_indexing_threads: 4 - pulsar_url: "pulsar://127.0.0.1:6650" - pulsar_tenant: "default" - pulsar_namespace: "default" - kube_namespace: "chroma" assignment_policy: RendezvousHashing: hasher: Murmur3 memberlist_provider: CustomResource: - memberlist_name: "worker-memberlist" + kube_namespace: "chroma" + memberlist_name: "query-service-memberlist" queue_size: 100 - ingest: - queue_size: 10000 sysdb: Grpc: - host: "coordinator.chroma" + host: "sysdb.chroma" port: 50051 - segment_manager: - storage_path: "./tmp/segment_manager/" storage: S3: bucket: "chroma-storage" + log: + Grpc: + host: "logservice.chroma" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + +compaction_service: + my_ip: "10.244.0.9" + my_port: 50051 + assignment_policy: + RendezvousHashing: + hasher: Murmur3 + memberlist_provider: + CustomResource: + kube_namespace: "chroma" + memberlist_name: "compaction-service-memberlist" + queue_size: 100 + sysdb: + Grpc: + host: "sysdb.chroma" + port: 50051 + storage: + S3: + bucket: "chroma-storage" + log: + Grpc: + host: "logservice.chroma" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + compactor: + compaction_manager_queue_size: 1000 + max_concurrent_jobs: 100 + compaction_interval_sec: 60 diff --git a/rust/worker/src/assignment/assignment_policy.rs b/rust/worker/src/assignment/assignment_policy.rs index bde70b26625..66b300bb29d 100644 --- a/rust/worker/src/assignment/assignment_policy.rs +++ b/rust/worker/src/assignment/assignment_policy.rs @@ -1,7 +1,4 @@ -use crate::{ - config::{Configurable, WorkerConfig}, - errors::ChromaError, -}; +use crate::{config::Configurable, errors::ChromaError}; use super::{ config::{AssignmentPolicyConfig, HasherType}, @@ -20,14 +17,10 @@ Interfaces /// This trait mirrors the go and python versions of the assignment policy /// interface. /// # Methods -/// - assign: Assign a key to a topic. +/// - assign: Assign a key to a member. /// - get_members: Get the members that can be assigned to. /// - set_members: Set the members that can be assigned to. -/// # Notes -/// An assignment policy is not responsible for creating the topics it assigns to. -/// It is the responsibility of the caller to ensure that the topics exist. -/// An assignment policy must be Send. -pub(crate) trait AssignmentPolicy: Send { +pub(crate) trait AssignmentPolicy: Send + Sync { fn assign(&self, key: &str) -> Result; fn get_members(&self) -> Vec; fn set_members(&mut self, members: Vec); @@ -45,31 +38,20 @@ pub(crate) struct RendezvousHashingAssignmentPolicy { } impl RendezvousHashingAssignmentPolicy { - // Rust beginners note - // The reason we take String and not &str is because we need to put the strings into our - // struct, and we can't do that with references so rather than clone the strings, we just - // take ownership of them and put the responsibility on the caller to clone them if they - // need to. This is the general pattern we should follow in rust - put the burden of cloning - // on the caller, and if they don't need to clone, they can pass ownership. - pub(crate) fn new( - pulsar_tenant: String, - pulsar_namespace: String, - ) -> RendezvousHashingAssignmentPolicy { + pub(crate) fn new() -> RendezvousHashingAssignmentPolicy { return RendezvousHashingAssignmentPolicy { hasher: Murmur3Hasher {}, members: vec![], }; } - - pub(crate) fn set_members(&mut self, members: Vec) { - self.members = members; - } } #[async_trait] -impl Configurable for RendezvousHashingAssignmentPolicy { - async fn try_from_config(worker_config: &WorkerConfig) -> Result> { - let assignment_policy_config = match &worker_config.assignment_policy { +impl Configurable for RendezvousHashingAssignmentPolicy { + async fn try_from_config( + config: &AssignmentPolicyConfig, + ) -> Result> { + let assignment_policy_config = match config { AssignmentPolicyConfig::RendezvousHashing(config) => config, }; let hasher = match assignment_policy_config.hasher { @@ -84,9 +66,8 @@ impl Configurable for RendezvousHashingAssignmentPolicy { impl AssignmentPolicy for RendezvousHashingAssignmentPolicy { fn assign(&self, key: &str) -> Result { - let topics = self.get_members(); - let topic = assign(key, topics, &self.hasher); - return topic; + let members = self.get_members(); + assign(key, members, &self.hasher) } fn get_members(&self) -> Vec { diff --git a/rust/worker/src/assignment/mod.rs b/rust/worker/src/assignment/mod.rs index 7ed1525f0bc..70be4c966cd 100644 --- a/rust/worker/src/assignment/mod.rs +++ b/rust/worker/src/assignment/mod.rs @@ -1,3 +1,17 @@ pub(crate) mod assignment_policy; pub(crate) mod config; -mod rendezvous_hash; +pub(crate) mod rendezvous_hash; + +use crate::{config::Configurable, errors::ChromaError}; + +use self::{assignment_policy::AssignmentPolicy, config::AssignmentPolicyConfig}; + +pub(crate) async fn from_config( + config: &AssignmentPolicyConfig, +) -> Result, Box> { + match &config { + crate::assignment::config::AssignmentPolicyConfig::RendezvousHashing(_) => Ok(Box::new( + assignment_policy::RendezvousHashingAssignmentPolicy::try_from_config(config).await?, + )), + } +} diff --git a/rust/worker/src/bin/compaction_service.rs b/rust/worker/src/bin/compaction_service.rs new file mode 100644 index 00000000000..01364a463d8 --- /dev/null +++ b/rust/worker/src/bin/compaction_service.rs @@ -0,0 +1,6 @@ +use worker::compaction_service_entrypoint; + +#[tokio::main] +async fn main() { + compaction_service_entrypoint().await; +} diff --git a/rust/worker/src/bin/query_service.rs b/rust/worker/src/bin/query_service.rs new file mode 100644 index 00000000000..f3cfa4c8282 --- /dev/null +++ b/rust/worker/src/bin/query_service.rs @@ -0,0 +1,6 @@ +use worker::query_service_entrypoint; + +#[tokio::main] +async fn main() { + query_service_entrypoint().await; +} diff --git a/rust/worker/src/bin/worker.rs b/rust/worker/src/bin/worker.rs deleted file mode 100644 index 16428d244ff..00000000000 --- a/rust/worker/src/bin/worker.rs +++ /dev/null @@ -1,6 +0,0 @@ -use worker::worker_entrypoint; - -#[tokio::main] -async fn main() { - worker_entrypoint().await; -} diff --git a/rust/worker/src/blockstore/arrow_blockfile/block/delta.rs b/rust/worker/src/blockstore/arrow_blockfile/block/delta.rs new file mode 100644 index 00000000000..99b7260c88d --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/block/delta.rs @@ -0,0 +1,453 @@ +use super::{Block, BlockBuilderOptions, BlockData, BlockDataBuilder}; +use crate::blockstore::{ + arrow_blockfile::{blockfile::MAX_BLOCK_SIZE, provider::ArrowBlockProvider}, + types::{BlockfileKey, KeyType, Value, ValueType}, +}; +use arrow::util::bit_util; +use parking_lot::RwLock; +use std::{collections::BTreeMap, sync::Arc}; + +/// A block delta tracks a source block and represents the new state of a block. Blocks are +/// immutable, so when a write is made to a block, a new block is created with the new state. +/// A block delta is a temporary representation of the new state of a block. A block delta +/// can be converted to a block data, which is then used to create a new block. A block data +/// can be converted into a block delta for new writes. +/// # Methods +/// - can_add: checks if a key value pair can be added to the block delta and still be within the +/// max block size. +/// - add: adds a key value pair to the block delta. +/// - delete: deletes a key from the block delta. +/// - get_min_key: gets the minimum key in the block delta. +/// - get_size: gets the size of the block delta. +/// - split: splits the block delta into two block deltas. +#[derive(Clone)] +pub struct BlockDelta { + pub source_block: Arc, + inner: Arc>, +} + +impl BlockDelta { + /// Checks if a key value pair can be added to the block delta and still be within the + /// max block size. + pub fn can_add(&self, key: &BlockfileKey, value: &Value) -> bool { + let inner = self.inner.read(); + inner.can_add(key, value) + } + + /// Adds a key value pair to the block delta. + pub fn add(&self, key: BlockfileKey, value: Value) { + let mut inner = self.inner.write(); + inner.add(key, value); + } + + /// Deletes a key from the block delta. + pub fn delete(&self, key: BlockfileKey) { + let mut inner = self.inner.write(); + inner.delete(key); + } + + /// Gets the minimum key in the block delta. + pub fn get_min_key(&self) -> Option { + let inner = self.inner.read(); + let first_key = inner.new_data.keys().next(); + first_key.cloned() + } + + /// Gets the size of the block delta as it would be in a block. This includes + /// the size of the prefix, key, and value data and the size of the offsets + /// where applicable. The size is rounded up to the nearest 64 bytes as per + /// the arrow specification. When a block delta is converted into a block data + /// the same sizing is used to allocate the memory for the block data. + pub fn get_size(&self) -> usize { + let inner = self.inner.read(); + inner.get_size( + self.source_block.get_key_type(), + self.source_block.get_value_type(), + ) + } + + /// Splits the block delta into two block deltas. The split point is the last key + /// that pushes the block over the half size. + /// # Arguments + /// - provider: the arrow block provider to create the new block. + /// # Returns + /// A tuple containing the the key of the split point and the new block delta. + /// The new block delta contains all the key value pairs after, but not including the + /// split point. + /// # Panics + /// This function will panic if their is no split point found. This should never happen + /// as we should only call this function if can_add returns false. + pub fn split(&self, provider: &ArrowBlockProvider) -> (BlockfileKey, BlockDelta) { + let new_block = provider.create_block( + self.source_block.get_key_type(), + self.source_block.get_value_type(), + ); + let mut inner = self.inner.write(); + let (split_key, new_adds) = inner.split( + self.source_block.get_key_type(), + self.source_block.get_value_type(), + ); + ( + split_key, + BlockDelta { + source_block: new_block, + inner: Arc::new(RwLock::new(BlockDeltaInner { new_data: new_adds })), + }, + ) + } + + fn get_prefix_size(&self) -> usize { + let inner = self.inner.read(); + inner.get_prefix_size() + } + + fn get_key_size(&self) -> usize { + let inner = self.inner.read(); + inner.get_key_size() + } + + fn get_value_size(&self) -> usize { + let inner = self.inner.read(); + inner.get_value_size() + } + + fn get_value_count(&self) -> usize { + let inner = self.inner.read(); + inner.get_value_count() + } + + fn len(&self) -> usize { + let inner = self.inner.read(); + inner.new_data.len() + } +} + +struct BlockDeltaInner { + new_data: BTreeMap, +} + +impl BlockDeltaInner { + fn add(&mut self, key: BlockfileKey, value: Value) { + self.new_data.insert(key, value); + } + + fn delete(&mut self, key: BlockfileKey) { + if self.new_data.contains_key(&key) { + self.new_data.remove(&key); + } + } + + fn get_block_size( + &self, + item_count: usize, + prefix_size: usize, + key_size: usize, + value_size: usize, + key_type: KeyType, + value_type: ValueType, + ) -> usize { + let prefix_total_bytes = bit_util::round_upto_multiple_of_64(prefix_size); + let prefix_offset_bytes = bit_util::round_upto_multiple_of_64((item_count + 1) * 4); + + // https://docs.rs/arrow/latest/arrow/array/array/struct.GenericListArray.html + let key_total_bytes = bit_util::round_upto_multiple_of_64(key_size); + let key_offset_bytes = self.offset_size_for_key_type(item_count, key_type); + + let value_total_bytes = bit_util::round_upto_multiple_of_64(value_size); + let value_offset_bytes = self.offset_size_for_value_type(item_count, value_type); + + prefix_total_bytes + + prefix_offset_bytes + + key_total_bytes + + key_offset_bytes + + value_total_bytes + + value_offset_bytes + } + + fn get_size(&self, key_type: KeyType, value_type: ValueType) -> usize { + let prefix_data_size = self.get_prefix_size(); + let key_data_size = self.get_key_size(); + let value_data_size = self.get_value_size(); + + self.get_block_size( + self.new_data.len(), + prefix_data_size, + key_data_size, + value_data_size, + key_type, + value_type, + ) + } + + fn get_prefix_size(&self) -> usize { + self.new_data + .iter() + .fold(0, |acc, (key, _)| acc + key.get_prefix_size()) + } + + fn get_key_size(&self) -> usize { + self.new_data + .iter() + .fold(0, |acc, (key, _)| acc + key.key.get_size()) + } + + fn get_value_size(&self) -> usize { + self.new_data + .iter() + .fold(0, |acc, (_, value)| acc + value.get_size()) + } + + fn get_value_count(&self) -> usize { + self.new_data.iter().fold(0, |acc, (_, value)| match value { + Value::Int32ArrayValue(arr) => acc + arr.len(), + Value::StringValue(s) => acc + s.len(), + Value::RoaringBitmapValue(bitmap) => acc + bitmap.serialized_size(), + Value::UintValue(_) => acc + 1, + _ => unimplemented!("Value type not implemented"), + }) + } + + fn can_add(&self, key: &BlockfileKey, value: &Value) -> bool { + let additional_prefix_size = key.get_prefix_size(); + let additional_key_size = key.key.get_size(); + let additional_value_size = value.get_size(); + + let prefix_data_size = self.get_prefix_size() + additional_prefix_size; + let key_data_size = self.get_key_size() + additional_key_size; + let value_data_size = self.get_value_size() + additional_value_size; + + let prefix_offset_size = bit_util::round_upto_multiple_of_64((self.new_data.len() + 1) * 4); + let key_offset_size = self.offset_size_for_key_type(self.new_data.len(), key.into()); + let value_offset_size = self.offset_size_for_value_type(self.new_data.len(), value.into()); + + let prefix_total_bytes = + bit_util::round_upto_multiple_of_64(prefix_data_size) + prefix_offset_size; + let key_total_bytes = bit_util::round_upto_multiple_of_64(key_data_size) + key_offset_size; + let value_total_bytes = + bit_util::round_upto_multiple_of_64(value_data_size) + value_offset_size; + let total_future_size = prefix_total_bytes + key_total_bytes + value_total_bytes; + + if total_future_size > MAX_BLOCK_SIZE { + return false; + } + + total_future_size <= MAX_BLOCK_SIZE + } + + fn offset_size_for_value_type(&self, item_count: usize, value_type: ValueType) -> usize { + match value_type { + ValueType::Int32Array | ValueType::String | ValueType::RoaringBitmap => { + bit_util::round_upto_multiple_of_64((item_count + 1) * 4) + } + ValueType::Uint => 0, + _ => unimplemented!("Value type not implemented"), + } + } + + fn offset_size_for_key_type(&self, item_count: usize, key_type: KeyType) -> usize { + match key_type { + KeyType::String => bit_util::round_upto_multiple_of_64((item_count + 1) * 4), + KeyType::Float | KeyType::Uint => 0, + _ => unimplemented!("Key type not implemented"), + } + } + + /// Splits the block delta into two block deltas. The split point is the last key + /// that pushes the block over the half size. + /// # Arguments + /// - key_type: the key type of the block. + /// - value_type: the value type of the block. + /// # Returns + /// + fn split( + &mut self, + key_type: KeyType, + value_type: ValueType, + ) -> (BlockfileKey, BTreeMap) { + let half_size = MAX_BLOCK_SIZE / 2; + let mut running_prefix_size = 0; + let mut running_key_size = 0; + let mut running_value_size = 0; + let mut running_count = 0; + + // The split key will be the last key that pushes the block over the half size. Not the first key that pushes it over + let mut split_key = None; + let mut iter = self.new_data.iter(); + while let Some((key, value)) = iter.next() { + running_prefix_size += key.get_prefix_size(); + running_key_size += key.key.get_size(); + running_value_size += value.get_size(); + running_count += 1; + let current_size = self.get_block_size( + running_count, + running_prefix_size, + running_key_size, + running_value_size, + key_type, + value_type, + ); + if current_size > half_size { + let next = iter.next(); + match next { + Some((next_key, _)) => split_key = Some(next_key.clone()), + None => split_key = Some(key.clone()), + } + break; + } + } + + match &split_key { + // Note: Consider returning a Result instead of panicking + // This should never happen as we should only call this + // function if can_add returns false. But it may be worth making + // this compile time safe. + None => panic!("No split point found"), + Some(split_key) => { + let split_after = self.new_data.split_off(split_key); + return (split_key.clone(), split_after); + } + } + } +} + +impl TryFrom<&BlockDelta> for BlockData { + type Error = super::BlockDataBuildError; + + fn try_from(delta: &BlockDelta) -> Result { + let mut builder = BlockDataBuilder::new( + delta.source_block.get_key_type(), + delta.source_block.get_value_type(), + Some(BlockBuilderOptions::new( + delta.len(), + delta.get_prefix_size(), + delta.get_key_size(), + delta.get_value_count(), + delta.get_value_size(), + )), + ); + for (key, value) in delta.inner.read().new_data.iter() { + builder.add(key.clone(), value.clone()); + } + builder.build() + } +} + +impl From> for BlockDelta { + fn from(source_block: Arc) -> Self { + // Read the exising block and put it into adds. We only create these + // when we have a write to this block, so we don't care about the cost of + // reading the block. Since we know we will have to do that no matter what. + let mut adds = BTreeMap::new(); + let source_block_iter = source_block.iter(); + for (key, value) in source_block_iter { + adds.insert(key, value); + } + BlockDelta { + source_block, + inner: Arc::new(RwLock::new(BlockDeltaInner { new_data: adds })), + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::blockstore::types::{Key, KeyType, ValueType}; + use arrow::array::Int32Array; + use rand::{random, Rng}; + + #[test] + fn test_sizing_int_arr_val() { + let block_provider = ArrowBlockProvider::new(); + let block = block_provider.create_block(KeyType::String, ValueType::Int32Array); + let delta = BlockDelta::from(block.clone()); + + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("prefix".to_string(), Key::String(format!("key{}", i))); + let value_len: usize = rand::thread_rng().gen_range(1..100); + let mut new_vec = Vec::with_capacity(value_len); + for _ in 0..value_len { + new_vec.push(random::()); + } + delta.add(key, Value::Int32ArrayValue(Int32Array::from(new_vec))); + } + + let size = delta.get_size(); + let block_data = BlockData::try_from(&delta).unwrap(); + assert_eq!(size, block_data.get_size()); + } + + #[test] + fn test_sizing_string_val() { + let block_provider = ArrowBlockProvider::new(); + let block = block_provider.create_block(KeyType::String, ValueType::String); + let delta = BlockDelta::from(block.clone()); + + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("prefix".to_string(), Key::String(format!("key{}", i))); + let value = Value::StringValue(format!("value{}", i)); + delta.add(key, value); + } + let size = delta.get_size(); + let block_data = BlockData::try_from(&delta).unwrap(); + assert_eq!(size, block_data.get_size()); + } + + #[test] + fn test_sizing_int_key() { + let block_provider = ArrowBlockProvider::new(); + let block = block_provider.create_block(KeyType::Float, ValueType::String); + let delta = BlockDelta::from(block.clone()); + + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("prefix".to_string(), Key::Float(i as f32)); + let value = Value::StringValue(format!("value{}", i)); + delta.add(key, value); + } + + let size = delta.get_size(); + let block_data = BlockData::try_from(&delta).unwrap(); + assert_eq!(size, block_data.get_size()); + } + + #[test] + fn test_sizing_roaring_bitmap_val() { + let block_provider = ArrowBlockProvider::new(); + let block = block_provider.create_block(KeyType::String, ValueType::RoaringBitmap); + let delta = BlockDelta::from(block.clone()); + + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::String(format!("{:04}", i))); + let value = Value::RoaringBitmapValue(roaring::RoaringBitmap::from_iter( + (0..i).map(|x| x as u32), + )); + delta.add(key, value); + } + + let size = delta.get_size(); + let block_data = BlockData::try_from(&delta).unwrap(); + assert_eq!(size, block_data.get_size()); + } + + #[test] + fn test_sizing_uint_key_val() { + let block_provider = ArrowBlockProvider::new(); + let block = block_provider.create_block(KeyType::Uint, ValueType::Uint); + let delta = BlockDelta::from(block.clone()); + + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("prefix".to_string(), Key::Uint(i as u32)); + let value = Value::UintValue(i as u32); + delta.add(key, value); + } + + let size = delta.get_size(); + let block_data = BlockData::try_from(&delta).unwrap(); + assert_eq!(size, block_data.get_size()); + } +} diff --git a/rust/worker/src/blockstore/arrow_blockfile/block/iterator.rs b/rust/worker/src/blockstore/arrow_blockfile/block/iterator.rs new file mode 100644 index 00000000000..5a771ea4ed3 --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/block/iterator.rs @@ -0,0 +1,123 @@ +use super::types::Block; +use crate::blockstore::types::{BlockfileKey, Key, KeyType, Value, ValueType}; +use arrow::array::{Array, BooleanArray, Int32Array, ListArray, StringArray, UInt32Array}; + +/// An iterator over the contents of a block. +/// This is a simple wrapper around the Arrow array data that is stored in the block. +/// For now, it clones the data in the Block, since it is only used to populate BlockDeltas. +pub(super) struct BlockIterator { + block: Block, + index: usize, + key_type: KeyType, + value_type: ValueType, +} + +impl BlockIterator { + pub fn new(block: Block, key_type: KeyType, value_type: ValueType) -> Self { + Self { + block, + index: 0, + key_type, + value_type, + } + } +} + +impl Iterator for BlockIterator { + type Item = (BlockfileKey, Value); + + fn next(&mut self) -> Option { + let data = &self.block.inner.read().data; + if data.is_none() { + return None; + } + + // Arrow requires us to downcast the array to the specific type we want to work with. + // This is a bit awkward, but it's the way Arrow works to allow for dynamic typing. + // We match and return None if the downcast fails, since we can't continue without the correct type. + // In practice, this should never happen, since we control the types of the data we store in the block and + // maintain the invariant that the data is always of the correct type. + + let prefix = match data + .as_ref() + .unwrap() + .data + .column(0) + .as_any() + .downcast_ref::() + { + Some(prefix) => prefix.value(self.index).to_owned(), + None => return None, + }; + + let key = match data.as_ref() { + Some(data) => data.data.column(1), + None => return None, + }; + + let value = match data.as_ref() { + Some(data) => data.data.column(2), + None => return None, + }; + + if self.index >= prefix.len() { + return None; + } + + let key = match self.key_type { + KeyType::String => match key.as_any().downcast_ref::() { + Some(key) => Key::String(key.value(self.index).to_string()), + None => return None, + }, + KeyType::Float => match key.as_any().downcast_ref::() { + Some(key) => Key::Float(key.value(self.index) as f32), + None => return None, + }, + KeyType::Bool => match key.as_any().downcast_ref::() { + Some(key) => Key::Bool(key.value(self.index)), + None => return None, + }, + KeyType::Uint => match key.as_any().downcast_ref::() { + Some(key) => Key::Uint(key.value(self.index) as u32), + None => return None, + }, + }; + + let value = match self.value_type { + ValueType::Int32Array => match value.as_any().downcast_ref::() { + Some(value) => { + let value = match value + .value(self.index) + .as_any() + .downcast_ref::() + { + Some(value) => { + // An arrow array, if nested in a larger structure, when cloned may clone the entire larger buffer. + // This leads to a memory overhead and also breaks our sizing assumptions. In order to work around this, + // we have to manuallly create a new array and copy the data over rather than relying on clone. + + // Note that we use a vector here to avoid the overhead of the builder. The from() method for primitive + // types uses unsafe code to wrap the vecs underlying buffer in an arrow array. + + // There are more performant ways to do this, but this is the most straightforward. + + let mut new_vec = Vec::with_capacity(value.len()); + for i in 0..value.len() { + new_vec.push(value.value(i)); + } + let value = Int32Array::from(new_vec); + Value::Int32ArrayValue(value) + } + None => return None, + }; + value + } + None => return None, + }, + // TODO: Implement the rest of the value types + _ => unimplemented!(), + }; + self.index += 1; + Some((BlockfileKey::new(prefix, key), value)) + } +} diff --git a/rust/worker/src/blockstore/arrow_blockfile/block/mod.rs b/rust/worker/src/blockstore/arrow_blockfile/block/mod.rs new file mode 100644 index 00000000000..d867bd65040 --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/block/mod.rs @@ -0,0 +1,6 @@ +mod iterator; +mod types; + +pub(in crate::blockstore::arrow_blockfile) mod delta; +// Re-export types at the arrow_blockfile module level +pub(in crate::blockstore::arrow_blockfile) use types::*; diff --git a/rust/worker/src/blockstore/arrow_blockfile/block/types.rs b/rust/worker/src/blockstore/arrow_blockfile/block/types.rs new file mode 100644 index 00000000000..0ccc805dfdf --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/block/types.rs @@ -0,0 +1,671 @@ +use crate::blockstore::types::{BlockfileKey, Key, KeyType, Value, ValueType}; +use crate::errors::{ChromaError, ErrorCodes}; +use arrow::array::{ + BinaryArray, BinaryBuilder, BooleanArray, BooleanBuilder, Float32Array, Float32Builder, UInt32Array, UInt32Builder, +}; +use arrow::{ + array::{Array, Int32Array, Int32Builder, ListArray, ListBuilder, StringArray, StringBuilder}, + datatypes::{DataType, Field}, + record_batch::RecordBatch, +}; +use parking_lot::RwLock; +use std::io::Error; +use std::sync::Arc; +use thiserror::Error; +use uuid::Uuid; + +use super::delta::BlockDelta; +use super::iterator::BlockIterator; + +/// BlockState represents the state of a block in the blockstore. Conceptually, a block is immutable once the broarder system +/// has been made aware of its existence. New blocks may exist locally but are not considered part of the blockstore until they +/// are registered. +/// ## State transitions +/// The state of a block is as follows: +/// - Uninitialized: The block has been created but no data has been added +/// - Initialized: Data has been added to the block but it has not been committed +/// - Commited: The block has been committed and is ready to be registered. At this point the block is immutable +/// - Registered: The block has been registered and is now part of the blockstore +#[derive(Clone, Copy)] +pub enum BlockState { + Uninitialized, + Initialized, + Commited, + Registered, +} + +pub(super) struct Inner { + pub(super) id: Uuid, + pub(super) data: Option, + pub(super) state: BlockState, + pub(super) key_type: KeyType, + pub(super) value_type: ValueType, +} + +/// A block in a blockfile. A block is a sorted collection of data that is immutable once it has been committed. +/// Blocks are the fundamental unit of storage in the blockstore and are used to store data in the form of (key, value) pairs. +/// These pairs are stored in an Arrow record batch with the schema (prefix, key, value). +/// Blocks are created in an uninitialized state and are transitioned to an initialized state once data has been added. Once +/// committed, a block is immutable and cannot be modified. Blocks are registered with the blockstore once they have been +/// flushed. +/// +/// ### BlockData Notes +/// A Block holds BlockData via its Inner. Conceptually, the BlockData being loaded into memory is an optimization. The Block interface +/// could also support out of core operations where the BlockData is loaded from disk on demand. Currently we force operations to be in-core +/// but could expand to out-of-core in the future. +#[derive(Clone)] +pub struct Block { + pub(super) inner: Arc>, +} + +#[derive(Error, Debug)] +pub enum BlockError { + #[error("Invalid state transition")] + InvalidStateTransition, + #[error("Block data error")] + BlockDataError(#[from] BlockDataBuildError), +} + +impl ChromaError for BlockError { + fn code(&self) -> ErrorCodes { + match self { + BlockError::InvalidStateTransition => ErrorCodes::Internal, + BlockError::BlockDataError(e) => e.code(), + } + } +} + +impl Block { + pub fn new(id: Uuid, key_type: KeyType, value_type: ValueType) -> Self { + Self { + inner: Arc::new(RwLock::new(Inner { + id, + data: None, + state: BlockState::Uninitialized, + key_type, + value_type, + })), + } + } + + pub fn get(&self, query_key: &BlockfileKey) -> Option { + match &self.inner.read().data { + Some(data) => { + let prefix = data.data.column(0); + let key = data.data.column(1); + let value = data.data.column(2); + // TODO: This should be binary search + for i in 0..prefix.len() { + if prefix + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + == query_key.prefix + { + let key_matches = match &query_key.key { + Key::String(inner_key) => { + inner_key + == key.as_any().downcast_ref::().unwrap().value(i) + } + Key::Float(inner_key) => { + *inner_key + == key + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + } + Key::Bool(inner_key) => { + *inner_key + == key + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + } + Key::Uint(inner_key) => { + *inner_key + == key.as_any().downcast_ref::().unwrap().value(i) + as u32 + } + }; + if key_matches { + match self.get_value_type() { + ValueType::Int32Array => { + return Some(Value::Int32ArrayValue( + value + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + )) + } + ValueType::String => { + return Some(Value::StringValue( + value + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + .to_string(), + )) + } + ValueType::RoaringBitmap => { + let bytes = value + .as_any() + .downcast_ref::() + .unwrap() + .value(i); + let bitmap = roaring::RoaringBitmap::deserialize_from(bytes); + match bitmap { + Ok(bitmap) => { + return Some(Value::RoaringBitmapValue(bitmap)) + } + // TODO: log error + Err(_) => return None, + } + } + ValueType::Uint => { + return Some(Value::UintValue( + value + .as_any() + .downcast_ref::() + .unwrap() + .value(i), + )) + } + // TODO: Add support for other types + _ => unimplemented!(), + } + } + } + } + None + } + None => None, + } + } + + /// Returns the size of the block in bytes + pub fn get_size(&self) -> usize { + match &self.inner.read().data { + Some(data) => data.get_size(), + None => 0, + } + } + + /// Returns the number of items in the block + pub fn len(&self) -> usize { + match &self.inner.read().data { + Some(data) => data.data.column(0).len(), + None => 0, + } + } + + pub fn get_id(&self) -> Uuid { + self.inner.read().id + } + + pub fn get_key_type(&self) -> KeyType { + self.inner.read().key_type + } + + pub fn get_value_type(&self) -> ValueType { + self.inner.read().value_type + } + + pub fn get_state(&self) -> BlockState { + self.inner.read().state + } + + /// Marks a block as commited. A commited block is immutable and is eligbile to be flushed and registered. + pub fn commit(&self) -> Result<(), Box> { + let mut inner = self.inner.write(); + match inner.state { + BlockState::Uninitialized => Ok(()), + BlockState::Initialized => { + inner.state = BlockState::Commited; + Ok(()) + } + BlockState::Commited | BlockState::Registered => { + Err(Box::new(BlockError::InvalidStateTransition)) + } + } + } + + pub fn apply_delta(&self, delta: &BlockDelta) -> Result<(), Box> { + let data = match BlockData::try_from(delta) { + Ok(data) => data, + Err(e) => return Err(Box::new(BlockError::BlockDataError(e))), + }; + let mut inner = self.inner.write(); + match inner.state { + BlockState::Uninitialized => { + inner.data = Some(data); + inner.state = BlockState::Initialized; + Ok(()) + } + BlockState::Initialized => { + inner.data = Some(data); + inner.state = BlockState::Initialized; + Ok(()) + } + BlockState::Commited | BlockState::Registered => { + Err(Box::new(BlockError::InvalidStateTransition)) + } + } + } + + pub(super) fn iter(&self) -> BlockIterator { + BlockIterator::new( + self.clone(), + self.inner.read().key_type, + self.inner.read().value_type, + ) + } +} + +/// BlockData represents the data in a block. The data is stored in an Arrow record batch with the column schema (prefix, key, value). +/// These are stored in sorted order by prefix and key for efficient lookups. +#[derive(Clone)] +pub(super) struct BlockData { + pub(super) data: RecordBatch, +} + +impl BlockData { + pub(crate) fn new(data: RecordBatch) -> Self { + Self { data } + } + + /// Returns the size of the block in bytes + pub(crate) fn get_size(&self) -> usize { + let mut total_size = 0; + for column in self.data.columns() { + total_size += column.get_buffer_memory_size(); + } + total_size + } +} + +// ============== BlockDataBuilder ============== + +enum KeyBuilder { + StringBuilder(StringBuilder), + FloatBuilder(Float32Builder), + BoolBuilder(BooleanBuilder), + UintBuilder(UInt32Builder), +} + +enum ValueBuilder { + Int32ArrayValueBuilder(ListBuilder), + StringValueBuilder(StringBuilder), + RoaringBitmapBuilder(BinaryBuilder), + UintValueBuilder(UInt32Builder), +} + +/// BlockDataBuilder is used to build a block. It is used to add data to a block and then build the BlockData once all data has been added. +/// It is only used internally to an arrow_blockfile. +pub(super) struct BlockDataBuilder { + prefix_builder: StringBuilder, + key_builder: KeyBuilder, + value_builder: ValueBuilder, + last_key: Option, +} + +/// ## Options for the BlockDataBuilder +/// - item_count: The number of items in the block +/// - prefix_data_capacity: The required capacity for the prefix data. This will be rounded to the nearest 64 byte alignment by arrow. +/// - key_data_capacity: The required capacity for the key data. This will be rounded to the nearest 64 byte alignment by arrow. +/// - total_value_count: The total number of values in the block +/// - total_value_capacity: The required capacity for the value data. This will be rounded to the nearest 64 byte alignment by arrow. +/// # Note +/// The capacities are the size of the initially allocated buffers. The builder will not enforce these limits and will grow the buffers as needed. +/// When in use in a blockfile, maintaining the block size is accomplished at the blockfile level. +pub(super) struct BlockBuilderOptions { + pub(super) item_count: usize, + pub(super) prefix_data_capacity: usize, + pub(super) key_data_capacity: usize, + pub(super) total_value_count: usize, + pub(super) total_value_capacity: usize, +} + +impl BlockBuilderOptions { + pub(super) fn new( + item_count: usize, + prefix_data_capacity: usize, + key_data_capacity: usize, + total_value_count: usize, + total_value_capacity: usize, + ) -> Self { + Self { + item_count, + prefix_data_capacity, + key_data_capacity, + total_value_count, + total_value_capacity, + } + } + + pub(super) fn default() -> Self { + Self { + item_count: 1024, + prefix_data_capacity: 1024, + key_data_capacity: 1024, + total_value_count: 1024, + total_value_capacity: 1024, + } + } +} + +impl BlockDataBuilder { + pub(super) fn new( + key_type: KeyType, + value_type: ValueType, + options: Option, + ) -> Self { + let options = options.unwrap_or(BlockBuilderOptions::default()); + let prefix_builder = + StringBuilder::with_capacity(options.item_count, options.prefix_data_capacity); + let key_builder = match key_type { + KeyType::String => KeyBuilder::StringBuilder(StringBuilder::with_capacity( + options.item_count, + options.key_data_capacity, + )), + KeyType::Float => { + KeyBuilder::FloatBuilder(Float32Builder::with_capacity(options.item_count)) + } + KeyType::Bool => { + KeyBuilder::BoolBuilder(BooleanBuilder::with_capacity(options.item_count)) + } + KeyType::Uint => { + KeyBuilder::UintBuilder(UInt32Builder::with_capacity(options.item_count)) + } + }; + let value_builder = match value_type { + ValueType::Int32Array => { + ValueBuilder::Int32ArrayValueBuilder(ListBuilder::with_capacity( + Int32Builder::with_capacity(options.total_value_count), + options.item_count, + )) + } + ValueType::String => ValueBuilder::StringValueBuilder(StringBuilder::with_capacity( + options.item_count, + options.total_value_capacity, + )), + ValueType::Uint => { + ValueBuilder::UintValueBuilder(UInt32Builder::with_capacity(options.item_count)) + } + ValueType::RoaringBitmap => ValueBuilder::RoaringBitmapBuilder( + BinaryBuilder::with_capacity(options.item_count, options.total_value_capacity), + ), + // TODO: Implement the other value types + _ => unimplemented!(), + }; + Self { + prefix_builder, + key_builder, + value_builder, + last_key: None, + } + } + + /// Adds a key, value pair to the block. The key must be greater than the last key added to the block otherwise an error is returned. + pub(super) fn add( + &mut self, + key: BlockfileKey, + value: Value, + ) -> Result<(), Box> { + match &self.last_key { + Some(last_key) => { + if key < *last_key { + return Err(Box::new(BlockDataAddError::KeyNotInOrder)); + } + } + None => {} + } + self.last_key = Some(key.clone()); + self.prefix_builder.append_value(key.prefix); + match self.key_builder { + KeyBuilder::StringBuilder(ref mut builder) => match key.key { + Key::String(key) => { + builder.append_value(key); + } + _ => unreachable!("Invalid key type for block"), + }, + KeyBuilder::FloatBuilder(ref mut builder) => match key.key { + Key::Float(key) => { + builder.append_value(key); + } + _ => unreachable!("Invalid key type for block"), + }, + KeyBuilder::BoolBuilder(ref mut builder) => match key.key { + Key::Bool(key) => { + builder.append_value(key); + } + _ => unreachable!("Invalid key type for block"), + }, + KeyBuilder::UintBuilder(ref mut builder) => match key.key { + Key::Uint(key) => { + builder.append_value(key); + } + _ => unreachable!("Invalid key type for block"), + }, + } + + match self.value_builder { + ValueBuilder::Int32ArrayValueBuilder(ref mut builder) => match value { + Value::Int32ArrayValue(array) => { + builder.append_value(&array); + } + _ => unreachable!("Invalid value type for block"), + }, + ValueBuilder::StringValueBuilder(ref mut builder) => match value { + Value::StringValue(string) => { + builder.append_value(string); + } + _ => unreachable!("Invalid value type for block"), + }, + ValueBuilder::UintValueBuilder(ref mut builder) => match value { + Value::UintValue(uint) => { + builder.append_value(uint); + } + _ => unreachable!("Invalid value type for block"), + }, + ValueBuilder::RoaringBitmapBuilder(ref mut builder) => match value { + Value::RoaringBitmapValue(bitmap) => { + let mut bytes = Vec::with_capacity(bitmap.serialized_size()); + match bitmap.serialize_into(&mut bytes) { + Ok(_) => builder.append_value(&bytes), + Err(e) => { + return Err(Box::new(BlockDataAddError::RoaringBitmapError(e))); + } + } + } + _ => unreachable!("Invalid value type for block"), + }, + } + + Ok(()) + } + + pub(super) fn build(&mut self) -> Result { + let prefix = self.prefix_builder.finish(); + let prefix_field = Field::new("prefix", DataType::Utf8, true); + // TODO: figure out how to get rid of nullable, the builders turn it on by default but we don't want it + let key_field; + let key = match self.key_builder { + KeyBuilder::StringBuilder(ref mut builder) => { + key_field = Field::new("key", DataType::Utf8, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + KeyBuilder::FloatBuilder(ref mut builder) => { + key_field = Field::new("key", DataType::Float32, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + KeyBuilder::BoolBuilder(ref mut builder) => { + key_field = Field::new("key", DataType::Boolean, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + KeyBuilder::UintBuilder(ref mut builder) => { + key_field = Field::new("key", DataType::UInt32, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + }; + + let value_field; + let value = match self.value_builder { + ValueBuilder::Int32ArrayValueBuilder(ref mut builder) => { + value_field = Field::new( + "value", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + ValueBuilder::StringValueBuilder(ref mut builder) => { + value_field = Field::new("value", DataType::Utf8, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + ValueBuilder::UintValueBuilder(ref mut builder) => { + value_field = Field::new("value", DataType::UInt32, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + ValueBuilder::RoaringBitmapBuilder(ref mut builder) => { + value_field = Field::new("value", DataType::Binary, true); + let arr = builder.finish(); + (&arr as &dyn Array).slice(0, arr.len()) + } + }; + + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + prefix_field, + key_field, + value_field, + ])); + let record_batch = + RecordBatch::try_new(schema, vec![Arc::new(prefix), Arc::new(key), value]); + match record_batch { + Ok(record_batch) => Ok(BlockData::new(record_batch)), + Err(e) => Err(BlockDataBuildError::ArrowError(e)), + } + } +} + +#[derive(Error, Debug)] +pub enum BlockDataAddError { + #[error("Blockfile key not in order")] + KeyNotInOrder, + #[error("Roaring bitmap error")] + RoaringBitmapError(#[from] Error), +} + +impl ChromaError for BlockDataAddError { + fn code(&self) -> ErrorCodes { + match self { + BlockDataAddError::KeyNotInOrder => ErrorCodes::InvalidArgument, + BlockDataAddError::RoaringBitmapError(_) => ErrorCodes::Internal, + } + } +} + +#[derive(Error, Debug)] +pub enum BlockDataBuildError { + #[error("Arrow error")] + ArrowError(#[from] arrow::error::ArrowError), +} + +impl ChromaError for BlockDataBuildError { + fn code(&self) -> ErrorCodes { + match self { + BlockDataBuildError::ArrowError(_) => ErrorCodes::Internal, + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::blockstore::types::Key; + use arrow::array::Int32Array; + + #[test] + fn test_block_builder_can_add() { + let num_entries = 1000; + + let mut keys = Vec::new(); + let mut key_bytes = 0; + for i in 0..num_entries { + keys.push(Key::String(format!("{:04}", i))); + key_bytes += i.to_string().len(); + } + + let prefix = "key".to_string(); + let prefix_bytes = prefix.len() * num_entries; + let mut block_builder = BlockDataBuilder::new( + KeyType::String, + ValueType::Int32Array, + Some(BlockBuilderOptions::new( + num_entries, + prefix_bytes, + key_bytes, + num_entries, // 2 int32s per entry + num_entries * 2 * 4, // 2 int32s per entry + )), + ); + + for i in 0..num_entries { + block_builder + .add( + BlockfileKey::new(prefix.clone(), keys[i].clone()), + Value::Int32ArrayValue(Int32Array::from(vec![i as i32, (i + 1) as i32])), + ) + .unwrap(); + } + + // Basic sanity check + let block_data = block_builder.build().unwrap(); + assert_eq!(block_data.data.column(0).len(), num_entries); + assert_eq!(block_data.data.column(1).len(), num_entries); + assert_eq!(block_data.data.column(2).len(), num_entries); + } + + #[test] + fn test_out_of_order_key_fails() { + let mut block_builder = BlockDataBuilder::new( + KeyType::String, + ValueType::Int32Array, + Some(BlockBuilderOptions::default()), + ); + + block_builder + .add( + BlockfileKey::new("key".to_string(), Key::String("b".to_string())), + Value::Int32ArrayValue(Int32Array::from(vec![1, 2])), + ) + .unwrap(); + + let result = block_builder.add( + BlockfileKey::new("key".to_string(), Key::String("a".to_string())), + Value::Int32ArrayValue(Int32Array::from(vec![1, 2])), + ); + + match result { + Ok(_) => panic!("Expected error"), + Err(e) => { + assert_eq!(e.code(), ErrorCodes::InvalidArgument); + } + } + } +} diff --git a/rust/worker/src/blockstore/arrow_blockfile/blockfile.rs b/rust/worker/src/blockstore/arrow_blockfile/blockfile.rs new file mode 100644 index 00000000000..31af3679731 --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/blockfile.rs @@ -0,0 +1,608 @@ +use super::super::types::{Blockfile, BlockfileKey, Key, KeyType, Value, ValueType}; +use super::block::{BlockError, BlockState}; +use super::provider::ArrowBlockProvider; +use super::sparse_index::SparseIndex; +use crate::blockstore::arrow_blockfile::block::delta::BlockDelta; +use crate::blockstore::BlockfileError; +use crate::errors::ChromaError; +use parking_lot::Mutex; +use std::sync::Arc; +use thiserror::Error; +use uuid::Uuid; + +pub(super) const MAX_BLOCK_SIZE: usize = 16384; + +/// ArrowBlockfile is a blockfile implementation that uses Apache Arrow for the block storage. +/// It stores a sparse index over a set of blocks sorted by key. +/// It uses a block provider to create new blocks and to retrieve existing blocks. +#[derive(Clone)] +pub(crate) struct ArrowBlockfile { + key_type: KeyType, + value_type: ValueType, + block_provider: ArrowBlockProvider, + sparse_index: Arc>, + transaction_state: Option>, +} + +/// TransactionState is a helper struct to keep track of the state of a transaction. +/// It keeps a list of block deltas that are applied during the transaction and the new +/// sparse index that is created during the transaction. The sparse index is immutable +/// so we can replace the sparse index of the blockfile with the new sparse index after +/// the transaction is committed. +struct TransactionState { + block_delta: Mutex>, + sparse_index: Mutex>>>, +} + +impl TransactionState { + fn new() -> Self { + Self { + block_delta: Mutex::new(Vec::new()), + sparse_index: Mutex::new(None), + } + } + + /// Add a new block delta to the transaction state + fn add_delta(&self, delta: BlockDelta) { + let mut block_delta = self.block_delta.lock(); + block_delta.push(delta); + } + + /// Get the block delta for a specific block id + fn get_delta_for_block(&self, search_id: &Uuid) -> Option { + let block_delta = self.block_delta.lock(); + for delta in &*block_delta { + if delta.source_block.get_id() == *search_id { + return Some(delta.clone()); + } + } + None + } +} + +#[derive(Error, Debug)] +pub(crate) enum ArrowBlockfileError { + #[error("Block not found")] + BlockNotFoundError, + #[error("Block Error")] + BlockError(#[from] BlockError), + #[error("No split key found")] + NoSplitKeyFound, +} + +impl ChromaError for ArrowBlockfileError { + fn code(&self) -> crate::errors::ErrorCodes { + match self { + ArrowBlockfileError::BlockNotFoundError => crate::errors::ErrorCodes::NotFound, + ArrowBlockfileError::BlockError(err) => err.code(), + ArrowBlockfileError::NoSplitKeyFound => crate::errors::ErrorCodes::Internal, + } + } +} + +impl Blockfile for ArrowBlockfile { + fn get(&self, key: BlockfileKey) -> Result> { + let target_block_id = self.sparse_index.lock().get_target_block_id(&key); + let target_block = match self.block_provider.get_block(&target_block_id) { + None => return Err(Box::new(ArrowBlockfileError::BlockNotFoundError)), + Some(block) => block, + }; + let value = target_block.get(&key); + match value { + None => return Err(Box::new(BlockfileError::NotFoundError)), + Some(value) => Ok(value), + } + } + + fn get_by_prefix( + &self, + prefix: String, + ) -> Result, Box> { + unimplemented!(); + } + + fn get_gt( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + unimplemented!(); + } + + fn get_gte( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + unimplemented!(); + } + + fn get_lt( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + unimplemented!(); + } + + fn get_lte( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + unimplemented!(); + } + + fn set( + &mut self, + key: BlockfileKey, + value: Value, + ) -> Result<(), Box> { + // TODO: value must be smaller than the block size except for position lists, which are a special case + // where we split the value across multiple blocks + if !self.in_transaction() { + return Err(Box::new(BlockfileError::TransactionNotInProgress)); + } + + // Validate key type + match key.key { + Key::String(_) => { + if self.key_type != KeyType::String { + return Err(Box::new(BlockfileError::InvalidKeyType)); + } + } + Key::Float(_) => { + if self.key_type != KeyType::Float { + return Err(Box::new(BlockfileError::InvalidKeyType)); + } + } + Key::Bool(_) => { + if self.key_type != KeyType::Bool { + return Err(Box::new(BlockfileError::InvalidKeyType)); + } + } + Key::Uint(_) => { + if self.key_type != KeyType::Uint { + return Err(Box::new(BlockfileError::InvalidKeyType)); + } + } + } + + // Validate value type + match value { + Value::Int32ArrayValue(_) => { + if self.value_type != ValueType::Int32Array { + return Err(Box::new(BlockfileError::InvalidValueType)); + } + } + Value::StringValue(_) => { + if self.value_type != ValueType::String { + return Err(Box::new(BlockfileError::InvalidValueType)); + } + } + Value::IntValue(_) => { + if self.value_type != ValueType::Int { + return Err(Box::new(BlockfileError::InvalidValueType)); + } + } + Value::UintValue(_) => { + if self.value_type != ValueType::Uint { + return Err(Box::new(BlockfileError::InvalidValueType)); + } + } + Value::PositionalPostingListValue(_) => { + if self.value_type != ValueType::PositionalPostingList { + return Err(Box::new(BlockfileError::InvalidValueType)); + } + } + Value::RoaringBitmapValue(_) => { + if self.value_type != ValueType::RoaringBitmap { + return Err(Box::new(BlockfileError::InvalidValueType)); + } + } + } + + let transaction_state = match &self.transaction_state { + None => return Err(Box::new(BlockfileError::TransactionNotInProgress)), + Some(transaction_state) => transaction_state, + }; + + // Get the target block id for the key + let mut transaction_sparse_index = transaction_state.sparse_index.lock(); + let target_block_id = match *transaction_sparse_index { + None => self.sparse_index.lock().get_target_block_id(&key), + Some(ref index) => index.lock().get_target_block_id(&key), + }; + + // See if a delta for the target block already exists, if not create a new one and add it to the transaction state + // Creating a delta loads the block entirely into memory + let delta = match transaction_state.get_delta_for_block(&target_block_id) { + None => { + let target_block = match self.block_provider.get_block(&target_block_id) { + None => return Err(Box::new(ArrowBlockfileError::BlockNotFoundError)), + Some(block) => block, + }; + let delta = BlockDelta::from(target_block); + transaction_state.add_delta(delta.clone()); + delta + } + Some(delta) => delta, + }; + + // Check if we can add to the the delta without pushing the block over the max size. + // If we can't, we need to split the block and create a new delta + if delta.can_add(&key, &value) { + delta.add(key, value); + } else { + let (split_key, new_delta) = delta.split(&self.block_provider); + match *transaction_sparse_index { + None => { + let new_sparse_index = + Arc::new(Mutex::new(SparseIndex::from(&self.sparse_index.lock()))); + new_sparse_index + .lock() + .add_block(split_key, new_delta.source_block.get_id()); + *transaction_sparse_index = Some(new_sparse_index); + } + Some(ref index) => { + index + .lock() + .add_block(split_key, new_delta.source_block.get_id()); + } + } + transaction_state.add_delta(new_delta); + drop(transaction_sparse_index); + // Recursive call to add to the new appropriate delta + self.set(key, value)? + } + Ok(()) + } + + fn begin_transaction(&mut self) -> Result<(), Box> { + if self.in_transaction() { + return Err(Box::new(BlockfileError::TransactionInProgress)); + } + self.transaction_state = Some(Arc::new(TransactionState::new())); + Ok(()) + } + + fn commit_transaction(&mut self) -> Result<(), Box> { + if !self.in_transaction() { + return Err(Box::new(BlockfileError::TransactionNotInProgress)); + } + + let transaction_state = match self.transaction_state { + None => return Err(Box::new(BlockfileError::TransactionNotInProgress)), + Some(ref transaction_state) => transaction_state, + }; + + for delta in &*transaction_state.block_delta.lock() { + // Blocks are WORM, so if the block is uninitialized or initialized we can update it directly, if its registered, meaning the broader system is aware of it, + // we need to create a new block and update the sparse index to point to the new block + + match delta.source_block.get_state() { + BlockState::Uninitialized => { + match delta.source_block.apply_delta(&delta) { + Ok(_) => {} + Err(err) => { + return Err(Box::new(ArrowBlockfileError::BlockError(*err))); + } + } + match delta.source_block.commit() { + Ok(_) => {} + Err(err) => { + return Err(Box::new(ArrowBlockfileError::BlockError(*err))); + } + } + } + BlockState::Initialized => { + match delta.source_block.apply_delta(&delta) { + Ok(_) => {} + Err(err) => { + return Err(Box::new(ArrowBlockfileError::BlockError(*err))); + } + } + match delta.source_block.commit() { + Ok(_) => {} + Err(err) => { + return Err(Box::new(ArrowBlockfileError::BlockError(*err))); + } + } + } + BlockState::Commited | BlockState::Registered => { + // If the block is commited or registered, we need to create a new block and update the sparse index + let new_block = self + .block_provider + .create_block(self.key_type, self.value_type); + match new_block.apply_delta(&delta) { + Ok(_) => {} + Err(err) => { + return Err(Box::new(ArrowBlockfileError::BlockError(*err))); + } + } + let new_min_key = match delta.get_min_key() { + // This should never happen. We don't panic here because we want to return a proper error + None => return Err(Box::new(ArrowBlockfileError::NoSplitKeyFound)), + Some(key) => key, + }; + let mut transaction_sparse_index = transaction_state.sparse_index.lock(); + match *transaction_sparse_index { + None => { + let new_sparse_index = + Arc::new(Mutex::new(SparseIndex::from(&self.sparse_index.lock()))); + new_sparse_index.lock().replace_block( + delta.source_block.get_id(), + new_block.get_id(), + new_min_key, + ); + *transaction_sparse_index = Some(new_sparse_index); + } + Some(ref index) => { + index.lock().replace_block( + delta.source_block.get_id(), + new_block.get_id(), + new_min_key, + ); + } + } + match new_block.commit() { + Ok(_) => {} + Err(err) => { + return Err(Box::new(ArrowBlockfileError::BlockError(*err))); + } + } + } + } + } + + // update the sparse index + let mut transaction_state_sparse_index = transaction_state.sparse_index.lock(); + if transaction_state_sparse_index.is_some() { + self.sparse_index = transaction_state_sparse_index.take().unwrap(); + // unwrap is safe because we just checked it + } + + // Reset the transaction state + drop(transaction_state_sparse_index); + self.transaction_state = None; + Ok(()) + } +} + +impl ArrowBlockfile { + pub(super) fn new( + key_type: KeyType, + value_type: ValueType, + block_provider: ArrowBlockProvider, + ) -> Self { + let initial_block = block_provider.create_block(key_type.clone(), value_type.clone()); + Self { + sparse_index: Arc::new(Mutex::new(SparseIndex::new(initial_block.get_id()))), + transaction_state: None, + block_provider, + key_type, + value_type, + } + } + + fn in_transaction(&self) -> bool { + self.transaction_state.is_some() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + + #[test] + fn test_blockfile() { + let block_provider = ArrowBlockProvider::new(); + let mut blockfile = + ArrowBlockfile::new(KeyType::String, ValueType::Int32Array, block_provider); + + blockfile.begin_transaction().unwrap(); + let key1 = BlockfileKey::new("key".to_string(), Key::String("zzzz".to_string())); + blockfile + .set( + key1.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![1, 2, 3])), + ) + .unwrap(); + let key2 = BlockfileKey::new("key".to_string(), Key::String("aaaa".to_string())); + blockfile + .set( + key2, + Value::Int32ArrayValue(Int32Array::from(vec![4, 5, 6])), + ) + .unwrap(); + blockfile.commit_transaction().unwrap(); + + let value = blockfile.get(key1).unwrap(); + match value { + Value::Int32ArrayValue(array) => { + assert_eq!(array.values(), &[1, 2, 3]); + } + _ => panic!("Unexpected value type"), + } + } + + #[test] + fn test_splitting() { + let block_provider = ArrowBlockProvider::new(); + let mut blockfile = + ArrowBlockfile::new(KeyType::String, ValueType::Int32Array, block_provider); + + blockfile.begin_transaction().unwrap(); + let n = 1200; + for i in 0..n { + let string_key = format!("{:04}", i); + let key = BlockfileKey::new("key".to_string(), Key::String(string_key)); + blockfile + .set(key, Value::Int32ArrayValue(Int32Array::from(vec![i]))) + .unwrap(); + } + blockfile.commit_transaction().unwrap(); + + for i in 0..n { + let string_key = format!("{:04}", i); + let key = BlockfileKey::new("key".to_string(), Key::String(string_key)); + let res = blockfile.get(key).unwrap(); + match res { + Value::Int32ArrayValue(array) => { + assert_eq!(array.values(), &[i]); + } + _ => panic!("Unexpected value type"), + } + } + + // Sparse index should have 3 blocks + assert_eq!(blockfile.sparse_index.lock().len(), 3); + assert!(blockfile.sparse_index.lock().is_valid()); + + // Add 5 new entries to the first block + blockfile.begin_transaction().unwrap(); + for i in 0..5 { + let new_key = format! {"{:05}", i}; + let key = BlockfileKey::new("key".to_string(), Key::String(new_key)); + blockfile + .set(key, Value::Int32ArrayValue(Int32Array::from(vec![i]))) + .unwrap(); + } + blockfile.commit_transaction().unwrap(); + + // Sparse index should still have 3 blocks + assert_eq!(blockfile.sparse_index.lock().len(), 3); + assert!(blockfile.sparse_index.lock().is_valid()); + + // Add 1200 more entries, causing splits + blockfile.begin_transaction().unwrap(); + for i in n..n * 2 { + let new_key = format! {"{:04}", i}; + let key = BlockfileKey::new("key".to_string(), Key::String(new_key)); + blockfile + .set(key, Value::Int32ArrayValue(Int32Array::from(vec![i]))) + .unwrap(); + } + blockfile.commit_transaction().unwrap(); + } + + #[test] + fn test_string_value() { + let block_provider = ArrowBlockProvider::new(); + let mut blockfile = ArrowBlockfile::new(KeyType::String, ValueType::String, block_provider); + + blockfile.begin_transaction().unwrap(); + let n = 2000; + + for i in 0..n { + let string_key = format!("{:04}", i); + let key = BlockfileKey::new("key".to_string(), Key::String(string_key.clone())); + blockfile + .set(key, Value::StringValue(string_key.clone())) + .unwrap(); + } + blockfile.commit_transaction().unwrap(); + + for i in 0..n { + let string_key = format!("{:04}", i); + let key = BlockfileKey::new("key".to_string(), Key::String(string_key.clone())); + let res = blockfile.get(key).unwrap(); + match res { + Value::StringValue(string) => { + assert_eq!(string, string_key); + } + _ => panic!("Unexpected value type"), + } + } + } + + #[test] + fn test_int_key() { + let block_provider = ArrowBlockProvider::new(); + let mut blockfile = ArrowBlockfile::new(KeyType::Float, ValueType::String, block_provider); + + blockfile.begin_transaction().unwrap(); + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::Float(i as f32)); + blockfile + .set(key, Value::StringValue(format!("{:04}", i))) + .unwrap(); + } + blockfile.commit_transaction().unwrap(); + + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::Float(i as f32)); + let res = blockfile.get(key).unwrap(); + match res { + Value::StringValue(string) => { + assert_eq!(string, format!("{:04}", i)); + } + _ => panic!("Unexpected value type"), + } + } + } + + #[test] + fn test_roaring_bitmap_value() { + let block_provider = ArrowBlockProvider::new(); + let mut blockfile = + ArrowBlockfile::new(KeyType::String, ValueType::RoaringBitmap, block_provider); + + blockfile.begin_transaction().unwrap(); + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::String(format!("{:04}", i))); + blockfile + .set( + key, + Value::RoaringBitmapValue(roaring::RoaringBitmap::from_iter( + (0..i).map(|x| x as u32), + )), + ) + .unwrap(); + } + blockfile.commit_transaction().unwrap(); + + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::String(format!("{:04}", i))); + let res = blockfile.get(key).unwrap(); + match res { + Value::RoaringBitmapValue(bitmap) => { + assert_eq!(bitmap.len(), i as u64); + assert_eq!( + bitmap.iter().collect::>(), + (0..i).collect::>() + ); + } + _ => panic!("Unexpected value type"), + } + } + } + + #[test] + fn test_uint_key_val() { + let block_provider = ArrowBlockProvider::new(); + let mut blockfile = ArrowBlockfile::new(KeyType::Uint, ValueType::Uint, block_provider); + + blockfile.begin_transaction().unwrap(); + let n = 2000; + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::Uint(i as u32)); + blockfile.set(key, Value::UintValue(i as u32)).unwrap(); + } + blockfile.commit_transaction().unwrap(); + + for i in 0..n { + let key = BlockfileKey::new("key".to_string(), Key::Uint(i as u32)); + let res = blockfile.get(key).unwrap(); + match res { + Value::UintValue(val) => { + assert_eq!(val, i as u32); + } + _ => panic!("Unexpected value type"), + } + } + } +} diff --git a/rust/worker/src/blockstore/arrow_blockfile/mod.rs b/rust/worker/src/blockstore/arrow_blockfile/mod.rs new file mode 100644 index 00000000000..fdff38999eb --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/mod.rs @@ -0,0 +1,4 @@ +mod block; +mod blockfile; +mod provider; +mod sparse_index; diff --git a/rust/worker/src/blockstore/arrow_blockfile/provider.rs b/rust/worker/src/blockstore/arrow_blockfile/provider.rs new file mode 100644 index 00000000000..6b985b8e256 --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/provider.rs @@ -0,0 +1,89 @@ +use super::{block::Block, blockfile::ArrowBlockfile}; +use crate::blockstore::{ + provider::{BlockfileProvider, CreateError, OpenError}, + Blockfile, KeyType, ValueType, +}; +use parking_lot::RwLock; +use std::{collections::HashMap, sync::Arc}; +use uuid::Uuid; + +/// A BlockFileProvider that creates ArrowBlockfiles (Arrow-backed blockfiles used for production). +/// For now, it keeps a simple local cache of blockfiles. +pub(super) struct ArrowBlockfileProvider { + block_provider: ArrowBlockProvider, + files: HashMap>, +} + +impl BlockfileProvider for ArrowBlockfileProvider { + fn new() -> Self { + Self { + block_provider: ArrowBlockProvider::new(), + files: HashMap::new(), + } + } + + fn open(&self, path: &str) -> Result, Box> { + match self.files.get(path) { + Some(file) => Ok(file.clone()), + None => Err(Box::new(OpenError::NotFound)), + } + } + + fn create( + &mut self, + path: &str, + key_type: KeyType, + value_type: ValueType, + ) -> Result, Box> { + match self.files.get(path) { + Some(_) => Err(Box::new(CreateError::AlreadyExists)), + None => { + let blockfile = Box::new(ArrowBlockfile::new( + key_type, + value_type, + self.block_provider.clone(), + )); + self.files.insert(path.to_string(), blockfile); + Ok(self.files.get(path).unwrap().clone()) + } + } + } +} + +/// A simple local cache of Arrow-backed blocks, the blockfile provider passes this +/// to the ArrowBlockfile when it creates a new blockfile. So that the blockfile can manage and access blocks +/// # Note +/// The implementation is currently very simple and not intended for robust production use. We should +/// introduce a more sophisticated cache that can handle tiered eviction and other features. This interface +/// is a placeholder for that. +struct ArrowBlockProviderInner { + blocks: HashMap>, +} + +#[derive(Clone)] +pub(super) struct ArrowBlockProvider { + inner: Arc>, +} + +impl ArrowBlockProvider { + pub(super) fn new() -> Self { + Self { + inner: Arc::new(RwLock::new(ArrowBlockProviderInner { + blocks: HashMap::new(), + })), + } + } + + pub(super) fn create_block(&self, key_type: KeyType, value_type: ValueType) -> Arc { + let block = Arc::new(Block::new(Uuid::new_v4(), key_type, value_type)); + self.inner + .write() + .blocks + .insert(block.get_id(), block.clone()); + block + } + + pub(super) fn get_block(&self, id: &Uuid) -> Option> { + self.inner.read().blocks.get(id).cloned() + } +} diff --git a/rust/worker/src/blockstore/arrow_blockfile/sparse_index.rs b/rust/worker/src/blockstore/arrow_blockfile/sparse_index.rs new file mode 100644 index 00000000000..d99afc64498 --- /dev/null +++ b/rust/worker/src/blockstore/arrow_blockfile/sparse_index.rs @@ -0,0 +1,213 @@ +use crate::blockstore::types::BlockfileKey; +use std::collections::{BTreeMap, HashMap}; +use std::fmt::Debug; +use uuid::Uuid; + +/// A sentinel blockfilekey wrapper to represent the start blocks range +/// # Note +/// The start key is used to represent the first block in the sparse index, this makes +/// it easier to handle the case where the first block is split into two and also makes +/// determining the target block for a given key easier +#[derive(Clone, Debug)] +enum SparseIndexDelimiter { + Start, + Key(BlockfileKey), +} + +impl PartialEq for SparseIndexDelimiter { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (SparseIndexDelimiter::Start, SparseIndexDelimiter::Start) => true, + (SparseIndexDelimiter::Key(k1), SparseIndexDelimiter::Key(k2)) => k1 == k2, + _ => false, + } + } +} + +impl Eq for SparseIndexDelimiter {} + +impl PartialOrd for SparseIndexDelimiter { + fn partial_cmp(&self, other: &Self) -> Option { + match (self, other) { + (SparseIndexDelimiter::Start, SparseIndexDelimiter::Start) => { + Some(std::cmp::Ordering::Equal) + } + (SparseIndexDelimiter::Start, SparseIndexDelimiter::Key(_)) => { + Some(std::cmp::Ordering::Less) + } + (SparseIndexDelimiter::Key(_), SparseIndexDelimiter::Start) => { + Some(std::cmp::Ordering::Greater) + } + (SparseIndexDelimiter::Key(k1), SparseIndexDelimiter::Key(k2)) => k1.partial_cmp(k2), + } + } +} + +impl Ord for SparseIndexDelimiter { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + match (self, other) { + (SparseIndexDelimiter::Start, SparseIndexDelimiter::Start) => std::cmp::Ordering::Equal, + (SparseIndexDelimiter::Start, SparseIndexDelimiter::Key(_)) => std::cmp::Ordering::Less, + (SparseIndexDelimiter::Key(_), SparseIndexDelimiter::Start) => { + std::cmp::Ordering::Greater + } + (SparseIndexDelimiter::Key(k1), SparseIndexDelimiter::Key(k2)) => k1.cmp(k2), + } + } +} + +/// A sparse index is used by a Blockfile to map a range of keys to a block id +/// # Methods +/// - `new` - Create a new sparse index with a single block +/// - `from` - Create a new sparse index from an existing sparse index +/// - `get_target_block_id` - Get the block id for a given key +/// - `add_block` - Add a new block to the sparse index +/// - `replace_block` - Replace an existing block with a new one +/// - `len` - Get the number of blocks in the sparse index +/// - `is_valid` - Check if the sparse index is valid, useful for debugging and testing +pub(super) struct SparseIndex { + forward: BTreeMap, + reverse: HashMap, +} + +impl SparseIndex { + pub(super) fn new(initial_block_id: Uuid) -> Self { + let mut forward = BTreeMap::new(); + forward.insert(SparseIndexDelimiter::Start, initial_block_id); + let mut reverse = HashMap::new(); + reverse.insert(initial_block_id, SparseIndexDelimiter::Start); + Self { forward, reverse } + } + + pub(super) fn from(old_sparse_index: &SparseIndex) -> Self { + Self { + forward: old_sparse_index.forward.clone(), + reverse: old_sparse_index.reverse.clone(), + } + } + + pub(super) fn get_target_block_id(&self, search_key: &BlockfileKey) -> Uuid { + let mut iter_curr = self.forward.iter(); + let mut iter_next = self.forward.iter().skip(1); + let search_key = SparseIndexDelimiter::Key(search_key.clone()); + while let Some((curr_key, curr_block_id)) = iter_curr.next() { + if let Some((next_key, _)) = iter_next.next() { + if search_key >= *curr_key && search_key < *next_key { + return *curr_block_id; + } + } else { + return *curr_block_id; + } + } + panic!("No blocks in the sparse index"); + } + + pub(super) fn add_block(&mut self, start_key: BlockfileKey, block_id: Uuid) { + self.forward + .insert(SparseIndexDelimiter::Key(start_key.clone()), block_id); + self.reverse + .insert(block_id, SparseIndexDelimiter::Key(start_key)); + } + + pub(super) fn replace_block( + &mut self, + old_block_id: Uuid, + new_block_id: Uuid, + new_start_key: BlockfileKey, + ) { + if let Some(old_start_key) = self.reverse.remove(&old_block_id) { + self.forward.remove(&old_start_key); + if old_start_key == SparseIndexDelimiter::Start { + self.forward + .insert(SparseIndexDelimiter::Start, new_block_id); + } else { + self.forward + .insert(SparseIndexDelimiter::Key(new_start_key), new_block_id); + } + } + } + + pub(super) fn len(&self) -> usize { + self.forward.len() + } + + /// Check if the sparse index is valid by ensuring that the keys are in order + pub(super) fn is_valid(&self) -> bool { + let mut first = true; + // Two pointer traversal to check if the keys are in order and that the start key is first + let mut iter_slow = self.forward.iter(); + let mut iter_fast = self.forward.iter().skip(1); + while let Some((curr_key, _)) = iter_slow.next() { + if first { + if curr_key != &SparseIndexDelimiter::Start { + return false; + } + first = false; + } + if let Some((next_key, _)) = iter_fast.next() { + if curr_key >= next_key { + return false; + } + } + } + true + } + + /// An iterator over the block uuids in the sparse index + pub(super) fn block_ids(&self) -> impl Iterator { + self.forward.values() + } +} + +impl Debug for SparseIndex { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "SparseIndex {{")?; + for (k, v) in self.forward.iter() { + write!(f, "\n {:?} -> {:?}", k, v)?; + } + write!(f, "\n}}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::blockstore::types::Key; + + #[test] + fn test_sparse_index() { + let block_id_1 = uuid::Uuid::new_v4(); + let mut sparse_index = SparseIndex::new(block_id_1); + let mut blockfile_key = + BlockfileKey::new("prefix".to_string(), Key::String("a".to_string())); + sparse_index.add_block(blockfile_key.clone(), block_id_1); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_1); + + blockfile_key = BlockfileKey::new("prefix".to_string(), Key::String("b".to_string())); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_1); + + // Split the range into two blocks (start, c), and (c, end) + let block_id_2 = uuid::Uuid::new_v4(); + blockfile_key = BlockfileKey::new("prefix".to_string(), Key::String("c".to_string())); + sparse_index.add_block(blockfile_key.clone(), block_id_2); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_2); + + // d should fall into the second block + blockfile_key = BlockfileKey::new("prefix".to_string(), Key::String("d".to_string())); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_2); + + // Split the second block into (c, f) and (f, end) + let block_id_3 = uuid::Uuid::new_v4(); + blockfile_key = BlockfileKey::new("prefix".to_string(), Key::String("f".to_string())); + sparse_index.add_block(blockfile_key.clone(), block_id_3); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_3); + + // g should fall into the third block + blockfile_key = BlockfileKey::new("prefix".to_string(), Key::String("g".to_string())); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_3); + + // b should fall into the first block + blockfile_key = BlockfileKey::new("prefix".to_string(), Key::String("b".to_string())); + assert_eq!(sparse_index.get_target_block_id(&blockfile_key), block_id_1); + } +} diff --git a/rust/worker/src/blockstore/mod.rs b/rust/worker/src/blockstore/mod.rs new file mode 100644 index 00000000000..4facc35f85d --- /dev/null +++ b/rust/worker/src/blockstore/mod.rs @@ -0,0 +1,8 @@ +mod arrow_blockfile; +mod positional_posting_list_value; +mod types; + +pub(crate) mod provider; + +pub(crate) use positional_posting_list_value::*; +pub(crate) use types::*; diff --git a/rust/worker/src/blockstore/positional_posting_list_value.rs b/rust/worker/src/blockstore/positional_posting_list_value.rs new file mode 100644 index 00000000000..0ca3ec91e54 --- /dev/null +++ b/rust/worker/src/blockstore/positional_posting_list_value.rs @@ -0,0 +1,233 @@ +use arrow::{ + array::{AsArray, Int32Array, Int32Builder, ListArray, ListBuilder}, + datatypes::Int32Type, +}; +use thiserror::Error; + +use std::collections::{HashMap, HashSet}; + +use crate::errors::{ChromaError, ErrorCodes}; + +#[derive(Debug, Clone)] +pub(crate) struct PositionalPostingList { + pub(crate) doc_ids: Int32Array, + pub(crate) positions: ListArray, +} + +impl PositionalPostingList { + pub(crate) fn get_doc_ids(&self) -> Int32Array { + return self.doc_ids.clone(); + } + + pub(crate) fn get_positions_for_doc_id(&self, doc_id: i32) -> Option { + let index = self.doc_ids.iter().position(|x| x == Some(doc_id)); + match index { + Some(index) => { + let target_positions = self.positions.value(index); + // Int32Array is composed of a Datatype, ScalarBuffer, and a null bitmap, these are all cheap to clone since the buffer is Arc'ed + let downcast = target_positions.as_primitive::().clone(); + return Some(downcast); + } + None => None, + } + } +} + +#[derive(Error, Debug)] +pub(crate) enum PositionalPostingListBuilderError { + #[error("Doc ID already exists in the list")] + DocIdAlreadyExists, + #[error("Doc ID does not exist in the list")] + DocIdDoesNotExist, + #[error("Incremental positions must be sorted")] + UnsortedPosition, +} + +impl ChromaError for PositionalPostingListBuilderError { + fn code(&self) -> ErrorCodes { + match self { + PositionalPostingListBuilderError::DocIdAlreadyExists => ErrorCodes::AlreadyExists, + PositionalPostingListBuilderError::DocIdDoesNotExist => ErrorCodes::InvalidArgument, + PositionalPostingListBuilderError::UnsortedPosition => ErrorCodes::InvalidArgument, + } + } +} + +pub(crate) struct PositionalPostingListBuilder { + doc_ids: HashSet, + positions: HashMap>, +} + +impl PositionalPostingListBuilder { + pub(crate) fn new() -> Self { + PositionalPostingListBuilder { + doc_ids: HashSet::new(), + positions: HashMap::new(), + } + } + + pub(crate) fn add_doc_id_and_positions( + &mut self, + doc_id: i32, + positions: Vec, + ) -> Result<(), PositionalPostingListBuilderError> { + if self.doc_ids.contains(&doc_id) { + return Err(PositionalPostingListBuilderError::DocIdAlreadyExists); + } + + self.doc_ids.insert(doc_id); + self.positions.insert(doc_id, positions); + Ok(()) + } + + pub(crate) fn contains_doc_id(&self, doc_id: i32) -> bool { + self.doc_ids.contains(&doc_id) + } + + pub(crate) fn add_positions_for_doc_id( + &mut self, + doc_id: i32, + positions: Vec, + ) -> Result<(), PositionalPostingListBuilderError> { + if !self.doc_ids.contains(&doc_id) { + return Err(PositionalPostingListBuilderError::DocIdDoesNotExist); + } + + self.positions.get_mut(&doc_id).unwrap().extend(positions); + Ok(()) + } + + pub(crate) fn build(&mut self) -> PositionalPostingList { + let mut doc_ids_builder = Int32Builder::new(); + let mut positions_builder = ListBuilder::new(Int32Builder::new()); + + let mut doc_ids_vec: Vec = self.doc_ids.drain().collect(); + doc_ids_vec.sort(); + let doc_ids_slice = doc_ids_vec.as_slice(); + doc_ids_builder.append_slice(doc_ids_slice); + let doc_ids = doc_ids_builder.finish(); + + for doc_id in doc_ids_slice.iter() { + // Get positions for the doc ID, sort them, put them into the positions_builder + let mut positions = self.positions.remove(doc_id).unwrap(); + positions.sort(); + let positions_as_some: Vec> = positions.into_iter().map(Some).collect(); + positions_builder.append_value(positions_as_some); + } + let positions = positions_builder.finish(); + + PositionalPostingList { + doc_ids: doc_ids, + positions: positions, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_positional_posting_list_single_document() { + let mut builder = PositionalPostingListBuilder::new(); + let _res = builder.add_doc_id_and_positions(1, vec![1, 2, 3]); + let list = builder.build(); + assert_eq!(list.get_doc_ids().values()[0], 1); + assert_eq!( + list.get_positions_for_doc_id(1).unwrap(), + Int32Array::from(vec![1, 2, 3]) + ); + } + + #[test] + fn test_positional_posting_list_multiple_documents() { + let mut builder = PositionalPostingListBuilder::new(); + let _res = builder.add_doc_id_and_positions(1, vec![1, 2, 3]); + let _res = builder.add_doc_id_and_positions(2, vec![4, 5, 6]); + let list = builder.build(); + assert_eq!(list.get_doc_ids().values()[0], 1); + assert_eq!(list.get_doc_ids().values()[1], 2); + assert_eq!( + list.get_positions_for_doc_id(1).unwrap(), + Int32Array::from(vec![1, 2, 3]) + ); + assert_eq!( + list.get_positions_for_doc_id(2).unwrap(), + Int32Array::from(vec![4, 5, 6]) + ); + } + + #[test] + fn test_positional_posting_list_document_ids_sorted_after_build() { + let mut builder = PositionalPostingListBuilder::new(); + let _res = builder.add_doc_id_and_positions(2, vec![4, 5, 6]); + let _res = builder.add_doc_id_and_positions(1, vec![1, 2, 3]); + let list = builder.build(); + assert_eq!(list.get_doc_ids().values()[0], 1); + assert_eq!(list.get_doc_ids().values()[1], 2); + assert_eq!( + list.get_positions_for_doc_id(1).unwrap(), + Int32Array::from(vec![1, 2, 3]) + ); + assert_eq!( + list.get_positions_for_doc_id(2).unwrap(), + Int32Array::from(vec![4, 5, 6]) + ); + } + + #[test] + fn test_positional_posting_list_all_positions_sorted_after_build() { + let mut builder = PositionalPostingListBuilder::new(); + let _res = builder.add_doc_id_and_positions(1, vec![3, 2, 1]); + let list = builder.build(); + assert_eq!(list.get_doc_ids().values()[0], 1); + assert_eq!( + list.get_positions_for_doc_id(1).unwrap(), + Int32Array::from(vec![1, 2, 3]) + ); + } + + #[test] + fn test_positional_posting_list_incremental_build() { + let mut builder = PositionalPostingListBuilder::new(); + + let _res = builder.add_doc_id_and_positions(1, vec![1, 2, 3]); + let _res = builder.add_positions_for_doc_id(1, [4].into()); + let _res = builder.add_positions_for_doc_id(1, [5].into()); + let _res = builder.add_positions_for_doc_id(1, [6].into()); + let _res = builder.add_doc_id_and_positions(2, vec![4, 5, 6]); + let _res = builder.add_positions_for_doc_id(2, [7].into()); + + let list = builder.build(); + assert_eq!(list.get_doc_ids().values()[0], 1); + assert_eq!(list.get_doc_ids().values()[1], 2); + assert_eq!( + list.get_positions_for_doc_id(1).unwrap(), + Int32Array::from(vec![1, 2, 3, 4, 5, 6]) + ); + } + + #[test] + fn test_all_positional_posting_list_behaviors_together() { + let mut builder = PositionalPostingListBuilder::new(); + + let _res = builder.add_doc_id_and_positions(1, vec![3, 2, 1]); + let _res = builder.add_positions_for_doc_id(1, [4].into()); + let _res = builder.add_positions_for_doc_id(1, [6].into()); + let _res = builder.add_positions_for_doc_id(1, [5].into()); + let _res = builder.add_doc_id_and_positions(2, vec![5, 4, 6]); + let _res = builder.add_positions_for_doc_id(2, [7].into()); + + let list = builder.build(); + assert_eq!(list.get_doc_ids().values()[0], 1); + assert_eq!(list.get_doc_ids().values()[1], 2); + assert_eq!( + list.get_positions_for_doc_id(1).unwrap(), + Int32Array::from(vec![1, 2, 3, 4, 5, 6]) + ); + assert_eq!( + list.get_positions_for_doc_id(2).unwrap(), + Int32Array::from(vec![4, 5, 6, 7]) + ); + } +} diff --git a/rust/worker/src/blockstore/provider.rs b/rust/worker/src/blockstore/provider.rs new file mode 100644 index 00000000000..5dbb433eb33 --- /dev/null +++ b/rust/worker/src/blockstore/provider.rs @@ -0,0 +1,98 @@ +use super::types::Blockfile; +use super::types::{HashMapBlockfile, KeyType, ValueType}; +use crate::errors::ChromaError; +use parking_lot::RwLock; +use std::collections::HashMap; +use std::sync::Arc; +use thiserror::Error; + +// =================== Interfaces =================== + +/// A trait for opening and creating blockfiles +/// # Methods +/// - new: Create a new instance of the blockfile provider. A blockfile provider returns a Box of a given type. +/// Currently, we support HashMap and Arrow-backed blockfiles. +/// - open: Open a blockfile at the given path, returning a Box and error if it does not exist +/// - create: Create a new blockfile at the given path, returning a Box and error if it already exists +/// # Example +/// ```ignore (TODO: This example is not runnable from outside the crate it seems. Fix this. Ignore for now.) +/// use crate::blockstore::provider::HashMapBlockfileProvider; +/// use crate::blockstore::types::{KeyType, ValueType}; +/// let mut provider = HashMapBlockfileProvider::new(); +/// let blockfile = provider.create("test", KeyType::String, ValueType::Int32Array); +/// ``` +pub(crate) trait BlockfileProvider { + fn new() -> Self; + fn open(&self, path: &str) -> Result, Box>; + fn create( + &mut self, + path: &str, + key_type: KeyType, + value_type: ValueType, + ) -> Result, Box>; +} + +/// A BlockFileProvider that creates HashMapBlockfiles (in-memory blockfiles used for testing). +/// It bookkeeps the blockfiles locally. +/// # Note +/// This is not intended for production use. +pub(crate) struct HashMapBlockfileProvider { + files: Arc>>>, +} + +impl BlockfileProvider for HashMapBlockfileProvider { + fn new() -> Self { + Self { + files: Arc::new(RwLock::new(HashMap::new())), + } + } + + fn open(&self, path: &str) -> Result, Box> { + match self.files.read().get(path) { + Some(file) => Ok(file.clone()), + None => Err(Box::new(OpenError::NotFound)), + } + } + + fn create( + &mut self, + path: &str, + key_type: KeyType, + value_type: ValueType, + ) -> Result, Box> { + let mut files = self.files.write(); + match files.get(path) { + Some(_) => Err(Box::new(CreateError::AlreadyExists)), + None => { + let blockfile = Box::new(HashMapBlockfile::new()); + files.insert(path.to_string(), blockfile); + Ok(files.get(path).unwrap().clone()) + } + } + } +} + +// =================== Errors =================== +#[derive(Error, Debug)] +pub(crate) enum OpenError { + #[error("Blockfile not found")] + NotFound, +} + +impl ChromaError for OpenError { + fn code(&self) -> crate::errors::ErrorCodes { + crate::errors::ErrorCodes::NotFound + } +} + +#[derive(Error, Debug)] +pub(crate) enum CreateError { + #[error("Blockfile already exists")] + AlreadyExists, +} + +impl ChromaError for CreateError { + fn code(&self) -> crate::errors::ErrorCodes { + crate::errors::ErrorCodes::AlreadyExists + } +} diff --git a/rust/worker/src/blockstore/types.rs b/rust/worker/src/blockstore/types.rs new file mode 100644 index 00000000000..8c7d9772807 --- /dev/null +++ b/rust/worker/src/blockstore/types.rs @@ -0,0 +1,628 @@ +use super::positional_posting_list_value::PositionalPostingList; +use crate::errors::{ChromaError, ErrorCodes}; +use arrow::array::{Array, Int32Array}; +use parking_lot::RwLock; +use roaring::RoaringBitmap; +use std::collections::HashMap; +use std::fmt::{Debug, Display}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; +use thiserror::Error; + +#[derive(Debug, Error)] +pub(crate) enum BlockfileError { + #[error("Key not found")] + NotFoundError, + #[error("Invalid Key Type")] + InvalidKeyType, + #[error("Invalid Value Type")] + InvalidValueType, + #[error("Transaction already in progress")] + TransactionInProgress, + #[error("Transaction not in progress")] + TransactionNotInProgress, +} + +impl ChromaError for BlockfileError { + fn code(&self) -> ErrorCodes { + match self { + BlockfileError::NotFoundError + | BlockfileError::InvalidKeyType + | BlockfileError::InvalidValueType => ErrorCodes::InvalidArgument, + BlockfileError::TransactionInProgress | BlockfileError::TransactionNotInProgress => { + ErrorCodes::FailedPrecondition + } + } + } +} + +// ===== Key Types ===== +#[derive(Clone)] +pub(crate) struct BlockfileKey { + pub(crate) prefix: String, + pub(crate) key: Key, +} + +impl Key { + pub(crate) fn get_size(&self) -> usize { + match self { + Key::String(s) => s.len(), + Key::Float(_) => 4, + Key::Bool(_) => 1, + Key::Uint(_) => 4, + } + } +} + +impl BlockfileKey { + pub(super) fn get_size(&self) -> usize { + self.get_prefix_size() + self.key.get_size() + } + + pub(super) fn get_prefix_size(&self) -> usize { + self.prefix.len() + } +} + +impl From<&BlockfileKey> for KeyType { + fn from(key: &BlockfileKey) -> Self { + match key.key { + Key::String(_) => KeyType::String, + Key::Float(_) => KeyType::Float, + Key::Bool(_) => KeyType::Bool, + Key::Uint(_) => KeyType::Uint, + } + } +} + +#[derive(Clone, PartialEq, PartialOrd, Debug)] +pub(crate) enum Key { + String(String), + Float(f32), + Bool(bool), + Uint(u32), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) enum KeyType { + String, + Float, + Bool, + Uint, +} + +impl Display for Key { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Key::String(s) => write!(f, "{}", s), + Key::Float(fl) => write!(f, "{}", fl), + Key::Bool(b) => write!(f, "{}", b), + Key::Uint(u) => write!(f, "{}", u), + } + } +} + +impl BlockfileKey { + pub(crate) fn new(prefix: String, key: Key) -> Self { + BlockfileKey { prefix, key } + } +} + +impl Debug for BlockfileKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "BlockfileKey(prefix: {}, key: {})", + self.prefix, self.key + ) + } +} + +impl Hash for BlockfileKey { + // Hash is only used for the HashMap implementation, which is a test/reference implementation + // Therefore this hash implementation is not used in production and allowed to be + // hacky + fn hash(&self, state: &mut H) { + self.prefix.hash(state); + } +} + +impl PartialEq for BlockfileKey { + fn eq(&self, other: &Self) -> bool { + self.prefix == other.prefix && self.key == other.key + } +} + +impl PartialOrd for BlockfileKey { + fn partial_cmp(&self, other: &Self) -> Option { + if self.prefix == other.prefix { + self.key.partial_cmp(&other.key) + } else { + self.prefix.partial_cmp(&other.prefix) + } + } +} + +impl Eq for BlockfileKey {} + +impl Ord for BlockfileKey { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + if self.prefix == other.prefix { + match self.key { + Key::String(ref s1) => match &other.key { + Key::String(s2) => s1.cmp(s2), + _ => panic!("Cannot compare string to float, bool, or uint"), + }, + Key::Float(f1) => match &other.key { + Key::Float(f2) => f1.partial_cmp(f2).unwrap(), + _ => panic!("Cannot compare float to string, bool, or uint"), + }, + Key::Bool(b1) => match &other.key { + Key::Bool(b2) => b1.cmp(b2), + _ => panic!("Cannot compare bool to string, float, or uint"), + }, + Key::Uint(u1) => match &other.key { + Key::Uint(u2) => u1.cmp(u2), + _ => panic!("Cannot compare uint to string, float, or bool"), + }, + } + } else { + self.prefix.cmp(&other.prefix) + } + } +} + +// ===== Value Types ===== + +#[derive(Debug)] +pub(crate) enum Value { + Int32ArrayValue(Int32Array), + PositionalPostingListValue(PositionalPostingList), + StringValue(String), + IntValue(i32), + UintValue(u32), + RoaringBitmapValue(RoaringBitmap), +} + +impl Clone for Value { + fn clone(&self) -> Self { + // TODO: make this correct for all types + match self { + Value::Int32ArrayValue(arr) => { + // An arrow array, if nested in a larger structure, when cloned may clone the entire larger buffer. + // This leads to a large memory overhead and also breaks our sizing assumptions. In order to work around this, + // we have to manuallly create a new array and copy the data over. + + // Note that we use a vector here to avoid the overhead of the builder. The from() method for primitive + // types uses unsafe code to wrap the vecs underlying buffer in an arrow array. + + // There are more performant ways to do this, but this is the most straightforward. + let mut new_vec = Vec::with_capacity(arr.len()); + for i in 0..arr.len() { + new_vec.push(arr.value(i)); + } + let new_arr = Int32Array::from(new_vec); + Value::Int32ArrayValue(new_arr) + } + Value::PositionalPostingListValue(list) => { + Value::PositionalPostingListValue(list.clone()) + } + Value::StringValue(s) => Value::StringValue(s.clone()), + Value::RoaringBitmapValue(bitmap) => Value::RoaringBitmapValue(bitmap.clone()), + Value::IntValue(i) => Value::IntValue(*i), + Value::UintValue(u) => Value::UintValue(*u), + } + } +} + +impl Value { + pub(crate) fn get_size(&self) -> usize { + match self { + Value::Int32ArrayValue(arr) => arr.get_buffer_memory_size(), + Value::PositionalPostingListValue(list) => { + unimplemented!("Size of positional posting list") + } + Value::StringValue(s) => s.len(), + Value::RoaringBitmapValue(bitmap) => bitmap.serialized_size(), + Value::IntValue(_) | Value::UintValue(_) => 4, + } + } +} + +impl From<&Value> for ValueType { + fn from(value: &Value) -> Self { + match value { + Value::Int32ArrayValue(_) => ValueType::Int32Array, + Value::PositionalPostingListValue(_) => ValueType::PositionalPostingList, + Value::RoaringBitmapValue(_) => ValueType::RoaringBitmap, + Value::StringValue(_) => ValueType::String, + Value::IntValue(_) => ValueType::Int, + Value::UintValue(_) => ValueType::Uint, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) enum ValueType { + Int32Array, + PositionalPostingList, + RoaringBitmap, + String, + Int, + Uint, +} + +pub(crate) trait Blockfile: BlockfileClone { + // ===== Transaction methods ===== + fn begin_transaction(&mut self) -> Result<(), Box>; + + fn commit_transaction(&mut self) -> Result<(), Box>; + + // ===== Data methods ===== + fn get(&self, key: BlockfileKey) -> Result>; + fn get_by_prefix( + &self, + prefix: String, + ) -> Result, Box>; + + fn set(&mut self, key: BlockfileKey, value: Value) -> Result<(), Box>; + + fn get_gt( + &self, + prefix: String, + key: Key, + ) -> Result, Box>; + + fn get_lt( + &self, + prefix: String, + key: Key, + ) -> Result, Box>; + + fn get_gte( + &self, + prefix: String, + key: Key, + ) -> Result, Box>; + + fn get_lte( + &self, + prefix: String, + key: Key, + ) -> Result, Box>; +} + +pub(crate) trait BlockfileClone { + fn clone_box(&self) -> Box; +} + +impl BlockfileClone for T +where + T: 'static + Blockfile + Clone, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + fn clone(&self) -> Box { + self.clone_box() + } +} + +#[derive(Clone)] +pub(crate) struct HashMapBlockfile { + map: Arc>>, +} + +impl HashMapBlockfile { + pub(super) fn new() -> Self { + Self { + map: Arc::new(RwLock::new(HashMap::new())), + } + } +} + +impl Blockfile for HashMapBlockfile { + fn get(&self, key: BlockfileKey) -> Result> { + match self.map.read().get(&key) { + Some(value) => Ok(value.clone()), + None => Err(Box::new(BlockfileError::NotFoundError)), + } + } + + fn get_by_prefix( + &self, + prefix: String, + ) -> Result, Box> { + let mut result = Vec::new(); + for (key, value) in self.map.read().iter() { + if key.prefix == prefix { + result.push((key.clone(), value.clone())); + } + } + Ok(result) + } + + fn set(&mut self, key: BlockfileKey, value: Value) -> Result<(), Box> { + self.map.write().insert(key, value); + Ok(()) + } + + fn get_gt( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + let mut result = Vec::new(); + for (k, v) in self.map.read().iter() { + if k.prefix == prefix && k.key > key { + result.push((k.clone(), v.clone())); + } + } + Ok(result) + } + + fn get_gte( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + let mut result = Vec::new(); + for (k, v) in self.map.read().iter() { + if k.prefix == prefix && k.key >= key { + result.push((k.clone(), v.clone())); + } + } + Ok(result) + } + + fn get_lt( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + let mut result = Vec::new(); + for (k, v) in self.map.read().iter() { + if k.prefix == prefix && k.key < key { + result.push((k.clone(), v.clone())); + } + } + Ok(result) + } + + fn get_lte( + &self, + prefix: String, + key: Key, + ) -> Result, Box> { + let mut result = Vec::new(); + for (k, v) in self.map.read().iter() { + if k.prefix == prefix && k.key <= key { + result.push((k.clone(), v.clone())); + } + } + Ok(result) + } + + fn begin_transaction(&mut self) -> Result<(), Box> { + Ok(()) + } + + fn commit_transaction(&mut self) -> Result<(), Box> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::blockstore::positional_posting_list_value::PositionalPostingListBuilder; + use arrow::array::Array; + + #[test] + fn test_blockfile_set_get() { + let mut blockfile = HashMapBlockfile::new(); + let key = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key1".to_string()), + }; + let _res = blockfile + .set( + key.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![1, 2, 3])), + ) + .unwrap(); + let value = blockfile.get(key); + // downcast to string + match value.unwrap() { + Value::Int32ArrayValue(arr) => assert_eq!(arr, Int32Array::from(vec![1, 2, 3])), + _ => panic!("Value is not a string"), + } + } + + #[test] + fn test_blockfile_get_by_prefix() { + let mut blockfile = HashMapBlockfile::new(); + let key1 = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key1".to_string()), + }; + let key2 = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key2".to_string()), + }; + let _res = blockfile + .set( + key1.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![1, 2, 3])), + ) + .unwrap(); + let _res = blockfile + .set( + key2.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![4, 5, 6])), + ) + .unwrap(); + let values = blockfile.get_by_prefix("text_prefix".to_string()).unwrap(); + assert_eq!(values.len(), 2); + // May return values in any order + match &values[0].1 { + Value::Int32ArrayValue(arr) => assert!( + arr == &Int32Array::from(vec![1, 2, 3]) || arr == &Int32Array::from(vec![4, 5, 6]) + ), + _ => panic!("Value is not a string"), + } + match &values[1].1 { + Value::Int32ArrayValue(arr) => assert!( + arr == &Int32Array::from(vec![1, 2, 3]) || arr == &Int32Array::from(vec![4, 5, 6]) + ), + _ => panic!("Value is not a string"), + } + } + + #[test] + fn test_bool_key() { + let mut blockfile = HashMapBlockfile::new(); + let key = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::Bool(true), + }; + let _res = blockfile.set( + key.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![1])), + ); + let value = blockfile.get(key).unwrap(); + match value { + Value::Int32ArrayValue(arr) => assert_eq!(arr, Int32Array::from(vec![1])), + _ => panic!("Value is not an arrow int32 array"), + } + } + + #[test] + fn test_storing_arrow_in_blockfile() { + let mut blockfile = HashMapBlockfile::new(); + let key = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key1".to_string()), + }; + let array = Value::Int32ArrayValue(Int32Array::from(vec![1, 2, 3])); + let _res = blockfile.set(key.clone(), array).unwrap(); + let value = blockfile.get(key).unwrap(); + match value { + Value::Int32ArrayValue(arr) => assert_eq!(arr, Int32Array::from(vec![1, 2, 3])), + _ => panic!("Value is not an arrow int32 array"), + } + } + + #[test] + fn test_blockfile_get_gt() { + let mut blockfile = HashMapBlockfile::new(); + let key1 = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key1".to_string()), + }; + let key2 = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key2".to_string()), + }; + let key3 = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("key3".to_string()), + }; + let _res = blockfile.set( + key1.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![1])), + ); + let _res = blockfile.set( + key2.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![2])), + ); + let _res = blockfile.set( + key3.clone(), + Value::Int32ArrayValue(Int32Array::from(vec![3])), + ); + let values = blockfile + .get_gt("text_prefix".to_string(), Key::String("key1".to_string())) + .unwrap(); + assert_eq!(values.len(), 2); + match &values[0].0.key { + Key::String(s) => assert!(s == "key2" || s == "key3"), + _ => panic!("Key is not a string"), + } + match &values[1].0.key { + Key::String(s) => assert!(s == "key2" || s == "key3"), + _ => panic!("Key is not a string"), + } + } + + #[test] + fn test_learning_arrow_struct() { + let mut builder = PositionalPostingListBuilder::new(); + let _res = builder.add_doc_id_and_positions(1, vec![0]); + let _res = builder.add_doc_id_and_positions(2, vec![0, 1]); + let _res = builder.add_doc_id_and_positions(3, vec![0, 1, 2]); + let list_term_1 = builder.build(); + + // Example of how to use the struct array, which is one value for a term + let mut blockfile = HashMapBlockfile::new(); + let key = BlockfileKey { + prefix: "text_prefix".to_string(), + key: Key::String("term1".to_string()), + }; + let _res = blockfile + .set(key.clone(), Value::PositionalPostingListValue(list_term_1)) + .unwrap(); + let posting_list = blockfile.get(key).unwrap(); + let posting_list = match posting_list { + Value::PositionalPostingListValue(arr) => arr, + _ => panic!("Value is not an arrow struct array"), + }; + + let ids = posting_list.get_doc_ids(); + let ids = ids.as_any().downcast_ref::().unwrap(); + // find index of target id + let target_id = 2; + + // imagine this is binary search instead of linear + for i in 0..ids.len() { + if ids.is_null(i) { + continue; + } + if ids.value(i) == target_id { + let pos_list = posting_list.get_positions_for_doc_id(target_id).unwrap(); + let pos_list = pos_list.as_any().downcast_ref::().unwrap(); + assert_eq!(pos_list.len(), 2); + assert_eq!(pos_list.value(0), 0); + assert_eq!(pos_list.value(1), 1); + break; + } + } + } + + #[test] + fn test_roaring_bitmap_example() { + let mut bitmap = RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(2); + bitmap.insert(3); + let mut blockfile = HashMapBlockfile::new(); + let key = BlockfileKey::new( + "text_prefix".to_string(), + Key::String("bitmap1".to_string()), + ); + let _res = blockfile + .set(key.clone(), Value::RoaringBitmapValue(bitmap)) + .unwrap(); + let value = blockfile.get(key).unwrap(); + match value { + Value::RoaringBitmapValue(bitmap) => { + assert!(bitmap.contains(1)); + assert!(bitmap.contains(2)); + assert!(bitmap.contains(3)); + } + _ => panic!("Value is not a roaring bitmap"), + } + } +} diff --git a/rust/worker/src/compactor/compaction_manager.rs b/rust/worker/src/compactor/compaction_manager.rs new file mode 100644 index 00000000000..2ca138bd2ad --- /dev/null +++ b/rust/worker/src/compactor/compaction_manager.rs @@ -0,0 +1,372 @@ +use super::scheduler::Scheduler; +use super::scheduler_policy::LasCompactionTimeSchedulerPolicy; +use crate::assignment::assignment_policy::AssignmentPolicy; +use crate::compactor::types::CompactionJob; +use crate::compactor::types::ScheduleMessage; +use crate::config::CompactionServiceConfig; +use crate::config::Configurable; +use crate::errors::ChromaError; +use crate::errors::ErrorCodes; +use crate::execution::operator::TaskMessage; +use crate::execution::orchestration::CompactOrchestrator; +use crate::execution::orchestration::CompactionResponse; +use crate::log::log::Log; +use crate::memberlist::Memberlist; +use crate::system::Component; +use crate::system::ComponentContext; +use crate::system::Handler; +use crate::system::Receiver; +use crate::system::System; +use async_trait::async_trait; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use std::fmt::Debug; +use std::fmt::Formatter; +use std::str::FromStr; +use std::time::Duration; +use thiserror::Error; +use uuid::Uuid; + +pub(crate) struct CompactionManager { + system: Option, + scheduler: Scheduler, + // Dependencies + log: Box, + // Dispatcher + dispatcher: Option>>, + // Config + compaction_manager_queue_size: usize, + compaction_interval: Duration, +} + +#[derive(Error, Debug)] +pub(crate) enum CompactionError { + #[error("Failed to compact")] + FailedToCompact, +} + +impl ChromaError for CompactionError { + fn code(&self) -> ErrorCodes { + match self { + CompactionError::FailedToCompact => ErrorCodes::Internal, + } + } +} + +impl CompactionManager { + pub(crate) fn new( + scheduler: Scheduler, + log: Box, + compaction_manager_queue_size: usize, + compaction_interval: Duration, + ) -> Self { + CompactionManager { + system: None, + scheduler, + log, + dispatcher: None, + compaction_manager_queue_size, + compaction_interval, + } + } + + async fn compact( + &self, + compaction_job: &CompactionJob, + ) -> Result> { + let collection_uuid = Uuid::from_str(&compaction_job.collection_id); + if collection_uuid.is_err() { + // handle error properly + println!("Failed to parse collection id"); + return Err(Box::new(CompactionError::FailedToCompact)); + } + + let dispatcher = match self.dispatcher { + Some(ref dispatcher) => dispatcher, + None => { + println!("No dispatcher found"); + return Err(Box::new(CompactionError::FailedToCompact)); + } + }; + + match self.system { + Some(ref system) => { + let orchestrator = CompactOrchestrator::new( + compaction_job.clone(), + system.clone(), + collection_uuid.unwrap(), + self.log.clone(), + dispatcher.clone(), + None, + ); + + match orchestrator.run().await { + Ok(result) => { + println!("Compaction Job completed"); + return Ok(result); + } + Err(e) => { + println!("Compaction Job failed"); + return Err(e); + } + } + } + None => { + println!("No system found"); + return Err(Box::new(CompactionError::FailedToCompact)); + } + }; + } + + // TODO: make the return type more informative + pub(crate) async fn compact_batch(&mut self) -> (u32, u32) { + self.scheduler.schedule().await; + let mut jobs = FuturesUnordered::new(); + for job in self.scheduler.get_jobs() { + jobs.push(self.compact(job)); + } + let mut num_completed_jobs = 0; + let mut num_failed_jobs = 0; + while let Some(job) = jobs.next().await { + match job { + Ok(result) => { + println!("Compaction completed: {:?}", result); + num_completed_jobs += 1; + } + Err(e) => { + println!("Compaction failed: {:?}", e); + num_failed_jobs += 1; + } + } + } + (num_completed_jobs, num_failed_jobs) + } + + pub(crate) fn set_dispatcher(&mut self, dispatcher: Box>) { + self.dispatcher = Some(dispatcher); + } + + pub(crate) fn set_system(&mut self, system: System) { + self.system = Some(system); + } +} + +#[async_trait] +impl Configurable for CompactionManager { + async fn try_from_config( + config: &crate::config::CompactionServiceConfig, + ) -> Result> { + let sysdb_config = &config.sysdb; + let sysdb = match crate::sysdb::from_config(sysdb_config).await { + Ok(sysdb) => sysdb, + Err(err) => { + return Err(err); + } + }; + let log_config = &config.log; + let log = match crate::log::from_config(log_config).await { + Ok(log) => log, + Err(err) => { + return Err(err); + } + }; + + let my_ip = config.my_ip.clone(); + let policy = Box::new(LasCompactionTimeSchedulerPolicy {}); + let compaction_interval_sec = config.compactor.compaction_interval_sec; + let max_concurrent_jobs = config.compactor.max_concurrent_jobs; + let compaction_manager_queue_size = config.compactor.compaction_manager_queue_size; + + let assignment_policy_config = &config.assignment_policy; + let assignment_policy = match crate::assignment::from_config(assignment_policy_config).await + { + Ok(assignment_policy) => assignment_policy, + Err(err) => { + return Err(err); + } + }; + let scheduler = Scheduler::new( + my_ip, + log.clone(), + sysdb.clone(), + policy, + max_concurrent_jobs, + assignment_policy, + ); + Ok(CompactionManager::new( + scheduler, + log, + compaction_manager_queue_size, + Duration::from_secs(compaction_interval_sec), + )) + } +} + +// ============== Component Implementation ============== +#[async_trait] +impl Component for CompactionManager { + fn queue_size(&self) -> usize { + self.compaction_manager_queue_size + } + + async fn on_start(&mut self, ctx: &crate::system::ComponentContext) -> () { + ctx.scheduler.schedule_interval( + ctx.sender.clone(), + ScheduleMessage {}, + self.compaction_interval, + None, + ctx, + ); + } +} + +impl Debug for CompactionManager { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CompactionManager") + } +} + +// ============== Handlers ============== +#[async_trait] +impl Handler for CompactionManager { + async fn handle( + &mut self, + _message: ScheduleMessage, + _ctx: &ComponentContext, + ) { + self.compact_batch().await; + } +} + +#[async_trait] +impl Handler for CompactionManager { + async fn handle(&mut self, message: Memberlist, _ctx: &ComponentContext) { + self.scheduler.set_memberlist(message); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::assignment::assignment_policy::RendezvousHashingAssignmentPolicy; + use crate::execution::dispatcher::Dispatcher; + use crate::log::log::InMemoryLog; + use crate::log::log::InternalLogRecord; + use crate::sysdb::test_sysdb::TestSysDb; + use crate::types::Collection; + use crate::types::LogRecord; + use crate::types::Operation; + use crate::types::OperationRecord; + use std::str::FromStr; + use uuid::Uuid; + + #[tokio::test] + async fn test_compaction_manager() { + let mut log = Box::new(InMemoryLog::new()); + + let collection_uuid_1 = Uuid::from_str("00000000-0000-0000-0000-000000000001").unwrap(); + let collection_id_1 = collection_uuid_1.to_string(); + log.add_log( + collection_id_1.clone(), + Box::new(InternalLogRecord { + collection_id: collection_id_1.clone(), + log_offset: 1, + log_ts: 1, + record: LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + }), + ); + + let collection_uuid_2 = Uuid::from_str("00000000-0000-0000-0000-000000000002").unwrap(); + let collection_id_2 = collection_uuid_2.to_string(); + log.add_log( + collection_id_2.clone(), + Box::new(InternalLogRecord { + collection_id: collection_id_2.clone(), + log_offset: 2, + log_ts: 2, + record: LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + }), + ); + + let mut sysdb = Box::new(TestSysDb::new()); + + let collection_1 = Collection { + id: collection_uuid_1, + name: "collection_1".to_string(), + metadata: None, + dimension: Some(1), + tenant: "tenant_1".to_string(), + database: "database_1".to_string(), + log_position: 0, + version: 0, + }; + + let collection_2 = Collection { + id: collection_uuid_2, + name: "collection_2".to_string(), + metadata: None, + dimension: Some(1), + tenant: "tenant_2".to_string(), + database: "database_2".to_string(), + log_position: 0, + version: 0, + }; + sysdb.add_collection(collection_1); + sysdb.add_collection(collection_2); + + let my_ip = "127.0.0.1".to_string(); + let compaction_manager_queue_size = 1000; + let max_concurrent_jobs = 10; + let compaction_interval = Duration::from_secs(1); + + // Set assignment policy + let mut assignment_policy = Box::new(RendezvousHashingAssignmentPolicy::new()); + assignment_policy.set_members(vec![my_ip.clone()]); + + let mut scheduler = Scheduler::new( + my_ip.clone(), + log.clone(), + sysdb.clone(), + Box::new(LasCompactionTimeSchedulerPolicy {}), + max_concurrent_jobs, + assignment_policy, + ); + // Set memberlist + scheduler.set_memberlist(vec![my_ip.clone()]); + + let mut manager = CompactionManager::new( + scheduler, + log, + compaction_manager_queue_size, + compaction_interval, + ); + + let system = System::new(); + + let dispatcher = Dispatcher::new(10, 10, 10); + let dispatcher_handle = system.start_component(dispatcher); + manager.set_dispatcher(dispatcher_handle.receiver()); + manager.set_system(system); + let (num_completed, number_failed) = manager.compact_batch().await; + assert_eq!(num_completed, 2); + assert_eq!(number_failed, 0); + } +} diff --git a/rust/worker/src/compactor/config.rs b/rust/worker/src/compactor/config.rs new file mode 100644 index 00000000000..729a2e2d1da --- /dev/null +++ b/rust/worker/src/compactor/config.rs @@ -0,0 +1,8 @@ +use serde::Deserialize; + +#[derive(Deserialize)] +pub(crate) struct CompactorConfig { + pub(crate) compaction_manager_queue_size: usize, + pub(crate) max_concurrent_jobs: usize, + pub(crate) compaction_interval_sec: u64, +} diff --git a/rust/worker/src/compactor/mod.rs b/rust/worker/src/compactor/mod.rs new file mode 100644 index 00000000000..08e07c63a14 --- /dev/null +++ b/rust/worker/src/compactor/mod.rs @@ -0,0 +1,8 @@ +mod compaction_manager; +pub(crate) mod config; +mod scheduler; +mod scheduler_policy; +mod types; + +pub(crate) use compaction_manager::*; +pub(crate) use types::*; diff --git a/rust/worker/src/compactor/scheduler.rs b/rust/worker/src/compactor/scheduler.rs new file mode 100644 index 00000000000..12dc1be433a --- /dev/null +++ b/rust/worker/src/compactor/scheduler.rs @@ -0,0 +1,287 @@ +use crate::assignment::assignment_policy::AssignmentPolicy; +use crate::compactor::scheduler_policy::SchedulerPolicy; +use crate::compactor::types::CompactionJob; +use crate::log::log::CollectionInfo; +use crate::log::log::CollectionRecord; +use crate::log::log::Log; +use crate::memberlist::Memberlist; +use crate::sysdb::sysdb::SysDb; +use uuid::Uuid; + +pub(crate) struct Scheduler { + my_ip: String, + log: Box, + sysdb: Box, + policy: Box, + job_queue: Vec, + max_concurrent_jobs: usize, + memberlist: Option, + assignment_policy: Box, +} + +impl Scheduler { + pub(crate) fn new( + my_ip: String, + log: Box, + sysdb: Box, + policy: Box, + max_concurrent_jobs: usize, + assignment_policy: Box, + ) -> Scheduler { + Scheduler { + my_ip, + log, + sysdb, + policy, + job_queue: Vec::with_capacity(max_concurrent_jobs), + max_concurrent_jobs, + memberlist: None, + assignment_policy, + } + } + + async fn get_collections_with_new_data(&mut self) -> Vec { + let collections = self.log.get_collections_with_new_data().await; + // TODO: filter collecitons based on memberlist + let collections = match collections { + Ok(collections) => collections, + Err(e) => { + // TODO: Log error + println!("Error: {:?}", e); + return Vec::new(); + } + }; + collections + } + + async fn verify_and_enrich_collections( + &mut self, + collections: Vec, + ) -> Vec { + let mut collection_records = Vec::new(); + for collection_info in collections { + let collection_id = Uuid::parse_str(collection_info.collection_id.as_str()); + if collection_id.is_err() { + // TODO: Log error + println!("Error: {:?}", collection_id.err()); + continue; + } + let collection_id = Some(collection_id.unwrap()); + // TODO: add a cache to avoid fetching the same collection multiple times + let result = self + .sysdb + .get_collections(collection_id, None, None, None) + .await; + + match result { + Ok(collection) => { + if collection.is_empty() { + // TODO: Log error + println!("Collection not found: {:?}", collection_info.collection_id); + continue; + } + collection_records.push(CollectionRecord { + id: collection[0].id.to_string(), + tenant_id: collection[0].tenant.clone(), + // TODO: get the last compaction time from the sysdb + last_compaction_time: 0, + first_record_time: collection_info.first_log_ts, + offset: collection_info.first_log_offset, + }); + } + Err(e) => { + // TODO: Log error + println!("Error: {:?}", e); + } + } + } + self.filter_collections(collection_records) + } + + fn filter_collections(&mut self, collections: Vec) -> Vec { + let mut filtered_collections = Vec::new(); + let members = self.memberlist.as_ref().unwrap(); + self.assignment_policy.set_members(members.clone()); + for collection in collections { + let result = self.assignment_policy.assign(collection.id.as_str()); + match result { + Ok(member) => { + if member == self.my_ip { + filtered_collections.push(collection); + } + } + Err(e) => { + // TODO: Log error + println!("Error: {:?}", e); + continue; + } + } + } + + filtered_collections + } + + pub(crate) async fn schedule_internal(&mut self, collection_records: Vec) { + let jobs = self + .policy + .determine(collection_records, self.max_concurrent_jobs as i32); + { + self.job_queue.clear(); + self.job_queue.extend(jobs); + } + } + + pub(crate) async fn schedule(&mut self) { + if self.memberlist.is_none() { + // TODO: Log error + println!("Memberlist is not set"); + return; + } + let collections = self.get_collections_with_new_data().await; + if collections.is_empty() { + return; + } + let collection_records = self.verify_and_enrich_collections(collections).await; + self.schedule_internal(collection_records).await; + } + + pub(crate) fn get_jobs(&self) -> impl Iterator { + self.job_queue.iter() + } + + pub(crate) fn set_memberlist(&mut self, memberlist: Memberlist) { + self.memberlist = Some(memberlist); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::assignment::assignment_policy::RendezvousHashingAssignmentPolicy; + use crate::compactor::scheduler_policy::LasCompactionTimeSchedulerPolicy; + use crate::log::log::InMemoryLog; + use crate::log::log::InternalLogRecord; + use crate::sysdb::test_sysdb::TestSysDb; + use crate::types::Collection; + use crate::types::LogRecord; + use crate::types::Operation; + use crate::types::OperationRecord; + use std::str::FromStr; + use uuid::Uuid; + + #[tokio::test] + async fn test_scheduler() { + let mut log = Box::new(InMemoryLog::new()); + + let collection_uuid_1 = Uuid::from_str("00000000-0000-0000-0000-000000000001").unwrap(); + let collection_id_1 = collection_uuid_1.to_string(); + log.add_log( + collection_id_1.clone(), + Box::new(InternalLogRecord { + collection_id: collection_id_1.clone(), + log_offset: 1, + log_ts: 1, + record: LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + }), + ); + + let collection_uuid_2 = Uuid::from_str("00000000-0000-0000-0000-000000000002").unwrap(); + let collection_id_2 = collection_uuid_2.to_string(); + log.add_log( + collection_id_2.clone(), + Box::new(InternalLogRecord { + collection_id: collection_id_2.clone(), + log_offset: 2, + log_ts: 2, + record: LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + }), + ); + + let mut sysdb = Box::new(TestSysDb::new()); + + let collection_1 = Collection { + id: collection_uuid_1, + name: "collection_1".to_string(), + metadata: None, + dimension: Some(1), + tenant: "tenant_1".to_string(), + database: "database_1".to_string(), + log_position: 0, + version: 0, + }; + + let collection_2 = Collection { + id: collection_uuid_2, + name: "collection_2".to_string(), + metadata: None, + dimension: Some(1), + tenant: "tenant_2".to_string(), + database: "database_2".to_string(), + log_position: 0, + version: 0, + }; + sysdb.add_collection(collection_1); + sysdb.add_collection(collection_2); + + let my_ip = "0.0.0.1".to_string(); + let scheduler_policy = Box::new(LasCompactionTimeSchedulerPolicy {}); + let max_concurrent_jobs = 1000; + + // Set assignment policy + let mut assignment_policy = Box::new(RendezvousHashingAssignmentPolicy::new()); + assignment_policy.set_members(vec![my_ip.clone()]); + + let mut scheduler = Scheduler::new( + my_ip.clone(), + log, + sysdb, + scheduler_policy, + max_concurrent_jobs, + assignment_policy, + ); + // Scheduler does nothing without memberlist + scheduler.schedule().await; + let jobs = scheduler.get_jobs(); + assert_eq!(jobs.count(), 0); + + // Set memberlist + scheduler.set_memberlist(vec![my_ip.clone()]); + scheduler.schedule().await; + let jobs = scheduler.get_jobs(); + + // TODO: 3/9 Tasks may be out of order since we have not yet implemented SysDB Get last compaction time. Use contains instead of equal. + let job_ids = jobs + .map(|t| t.collection_id.clone()) + .collect::>(); + assert_eq!(job_ids.len(), 2); + assert!(job_ids.contains(&collection_id_1)); + assert!(job_ids.contains(&collection_id_2)); + + // Test filter_collections + let member_1 = "0.0.0.1".to_string(); + let member_2 = "0.0.0.2".to_string(); + let members = vec![member_1.clone(), member_2.clone()]; + scheduler.set_memberlist(members.clone()); + scheduler.schedule().await; + let jobs = scheduler.get_jobs(); + assert_eq!(jobs.count(), 1); + } +} diff --git a/rust/worker/src/compactor/scheduler_policy.rs b/rust/worker/src/compactor/scheduler_policy.rs new file mode 100644 index 00000000000..9a3930d5de5 --- /dev/null +++ b/rust/worker/src/compactor/scheduler_policy.rs @@ -0,0 +1,88 @@ +use crate::compactor::types::CompactionJob; +use crate::log::log::CollectionRecord; + +pub(crate) trait SchedulerPolicy: Send + Sync + SchedulerPolicyClone { + fn determine(&self, collections: Vec, number_jobs: i32) + -> Vec; +} + +pub(crate) trait SchedulerPolicyClone { + fn clone_box(&self) -> Box; +} + +impl SchedulerPolicyClone for T +where + T: 'static + SchedulerPolicy + Clone, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + fn clone(&self) -> Box { + self.clone_box() + } +} + +#[derive(Clone)] +pub(crate) struct LasCompactionTimeSchedulerPolicy {} + +impl SchedulerPolicy for LasCompactionTimeSchedulerPolicy { + fn determine( + &self, + collections: Vec, + number_jobs: i32, + ) -> Vec { + let mut collections = collections; + collections.sort_by(|a, b| a.last_compaction_time.cmp(&b.last_compaction_time)); + let number_tasks = if number_jobs > collections.len() as i32 { + collections.len() as i32 + } else { + number_jobs + }; + let mut tasks = Vec::new(); + for collection in &collections[0..number_tasks as usize] { + tasks.push(CompactionJob { + collection_id: collection.id.clone(), + tenant_id: collection.tenant_id.clone(), + offset: collection.offset, + }); + } + tasks + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scheduler_policy() { + let scheduler_policy = LasCompactionTimeSchedulerPolicy {}; + let collections = vec![ + CollectionRecord { + id: "test1".to_string(), + tenant_id: "test".to_string(), + last_compaction_time: 1, + first_record_time: 1, + offset: 0, + }, + CollectionRecord { + id: "test2".to_string(), + tenant_id: "test".to_string(), + last_compaction_time: 0, + first_record_time: 0, + offset: 0, + }, + ]; + let jobs = scheduler_policy.determine(collections.clone(), 1); + assert_eq!(jobs.len(), 1); + assert_eq!(jobs[0].collection_id, "test2"); + + let jobs = scheduler_policy.determine(collections.clone(), 2); + assert_eq!(jobs.len(), 2); + assert_eq!(jobs[0].collection_id, "test2"); + assert_eq!(jobs[1].collection_id, "test1"); + } +} diff --git a/rust/worker/src/compactor/types.rs b/rust/worker/src/compactor/types.rs new file mode 100644 index 00000000000..f456e577414 --- /dev/null +++ b/rust/worker/src/compactor/types.rs @@ -0,0 +1,9 @@ +#[derive(Clone, Eq, PartialEq, Debug)] +pub(crate) struct CompactionJob { + pub(crate) collection_id: String, + pub(crate) tenant_id: String, + pub(crate) offset: i64, +} + +#[derive(Clone, Debug)] +pub(crate) struct ScheduleMessage {} diff --git a/rust/worker/src/config.rs b/rust/worker/src/config.rs index 7583bf0114e..e8f2a75f9df 100644 --- a/rust/worker/src/config.rs +++ b/rust/worker/src/config.rs @@ -19,7 +19,8 @@ const ENV_PREFIX: &str = "CHROMA_"; pub(crate) struct RootConfig { // The root config object wraps the worker config object so that // we can share the same config file between multiple services. - pub worker: WorkerConfig, + pub query_service: QueryServiceConfig, + pub compaction_service: CompactionServiceConfig, } impl RootConfig { @@ -85,30 +86,45 @@ impl RootConfig { /// # Description /// The primary config for the worker service. /// ## Description of parameters -/// - my_ip: The IP address of the worker service. Used for memberlist assignment. Must be provided -/// - num_indexing_threads: The number of indexing threads to use. If not provided, defaults to the number of cores on the machine. -/// - pulsar_tenant: The pulsar tenant to use. Must be provided. -/// - pulsar_namespace: The pulsar namespace to use. Must be provided. +/// - my_ip: The IP address of the worker service. Used for memberlist assignment. Must be provided. /// - assignment_policy: The assignment policy to use. Must be provided. /// # Notes /// In order to set the enviroment variables, you must prefix them with CHROMA_WORKER__. /// For example, to set my_ip, you would set CHROMA_WORKER__MY_IP. /// Each submodule that needs to be configured from the config object should implement the Configurable trait and /// have its own field in this struct for its Config struct. -pub(crate) struct WorkerConfig { +pub(crate) struct QueryServiceConfig { pub(crate) my_ip: String, pub(crate) my_port: u16, - pub(crate) num_indexing_threads: u32, - pub(crate) pulsar_tenant: String, - pub(crate) pulsar_namespace: String, - pub(crate) pulsar_url: String, - pub(crate) kube_namespace: String, pub(crate) assignment_policy: crate::assignment::config::AssignmentPolicyConfig, pub(crate) memberlist_provider: crate::memberlist::config::MemberlistProviderConfig, - pub(crate) ingest: crate::ingest::config::IngestConfig, pub(crate) sysdb: crate::sysdb::config::SysDbConfig, - pub(crate) segment_manager: crate::segment::config::SegmentManagerConfig, pub(crate) storage: crate::storage::config::StorageConfig, + pub(crate) log: crate::log::config::LogConfig, + pub(crate) dispatcher: crate::execution::config::DispatcherConfig, +} + +#[derive(Deserialize)] +/// # Description +/// The primary config for the compaction service. +/// ## Description of parameters +/// - my_ip: The IP address of the worker service. Used for memberlist assignment. Must be provided. +/// - assignment_policy: The assignment policy to use. Must be provided. +/// # Notes +/// In order to set the enviroment variables, you must prefix them with CHROMA_COMPACTOR__. +/// For example, to set my_ip, you would set CHROMA_COMPACTOR__MY_IP. +/// Each submodule that needs to be configured from the config object should implement the Configurable trait and +/// have its own field in this struct for its Config struct. +pub(crate) struct CompactionServiceConfig { + pub(crate) my_ip: String, + pub(crate) my_port: u16, + pub(crate) assignment_policy: crate::assignment::config::AssignmentPolicyConfig, + pub(crate) memberlist_provider: crate::memberlist::config::MemberlistProviderConfig, + pub(crate) sysdb: crate::sysdb::config::SysDbConfig, + pub(crate) storage: crate::storage::config::StorageConfig, + pub(crate) log: crate::log::config::LogConfig, + pub(crate) dispatcher: crate::execution::config::DispatcherConfig, + pub(crate) compactor: crate::compactor::config::CompactorConfig, } /// # Description @@ -117,8 +133,8 @@ pub(crate) struct WorkerConfig { /// This trait is used to configure structs from the config object. /// Components that need to be configured from the config object should implement this trait. #[async_trait] -pub(crate) trait Configurable { - async fn try_from_config(worker_config: &WorkerConfig) -> Result> +pub(crate) trait Configurable { + async fn try_from_config(worker_config: &T) -> Result> where Self: Sized; } @@ -134,40 +150,71 @@ mod tests { let _ = jail.create_file( "chroma_config.yaml", r#" - worker: + query_service: my_ip: "192.0.0.1" my_port: 50051 - num_indexing_threads: 4 - pulsar_tenant: "public" - pulsar_namespace: "default" - pulsar_url: "pulsar://localhost:6650" - kube_namespace: "chroma" assignment_policy: RendezvousHashing: hasher: Murmur3 memberlist_provider: CustomResource: - memberlist_name: "worker-memberlist" + kube_namespace: "chroma" + memberlist_name: "query-service-memberlist" queue_size: 100 - ingest: - queue_size: 100 sysdb: Grpc: host: "localhost" port: 50051 - segment_manager: - storage_path: "/tmp" storage: S3: bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + + compaction_service: + my_ip: "192.0.0.1" + my_port: 50051 + assignment_policy: + RendezvousHashing: + hasher: Murmur3 + memberlist_provider: + CustomResource: + kube_namespace: "chroma" + memberlist_name: "compaction-service-memberlist" + queue_size: 100 + sysdb: + Grpc: + host: "localhost" + port: 50051 + storage: + S3: + bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + compactor: + compaction_manager_queue_size: 1000 + max_concurrent_jobs: 100 + compaction_interval_sec: 60 "#, ); let config = RootConfig::load(); - assert_eq!(config.worker.my_ip, "192.0.0.1"); - assert_eq!(config.worker.num_indexing_threads, 4); - assert_eq!(config.worker.pulsar_tenant, "public"); - assert_eq!(config.worker.pulsar_namespace, "default"); - assert_eq!(config.worker.kube_namespace, "chroma"); + assert_eq!(config.query_service.my_ip, "192.0.0.1"); + assert_eq!(config.query_service.my_port, 50051); + + assert_eq!(config.compaction_service.my_ip, "192.0.0.1"); + assert_eq!(config.compaction_service.my_port, 50051); Ok(()) }); } @@ -178,41 +225,71 @@ mod tests { let _ = jail.create_file( "random_path.yaml", r#" - worker: + query_service: my_ip: "192.0.0.1" my_port: 50051 - num_indexing_threads: 4 - pulsar_tenant: "public" - pulsar_namespace: "default" - pulsar_url: "pulsar://localhost:6650" - kube_namespace: "chroma" assignment_policy: RendezvousHashing: hasher: Murmur3 memberlist_provider: CustomResource: - memberlist_name: "worker-memberlist" + kube_namespace: "chroma" + memberlist_name: "query-service-memberlist" queue_size: 100 - ingest: - queue_size: 100 sysdb: Grpc: host: "localhost" port: 50051 - segment_manager: - storage_path: "/tmp" storage: S3: bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + compaction_service: + my_ip: "192.0.0.1" + my_port: 50051 + assignment_policy: + RendezvousHashing: + hasher: Murmur3 + memberlist_provider: + CustomResource: + kube_namespace: "chroma" + memberlist_name: "compaction-service-memberlist" + queue_size: 100 + sysdb: + Grpc: + host: "localhost" + port: 50051 + storage: + S3: + bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + compactor: + compaction_manager_queue_size: 1000 + max_concurrent_jobs: 100 + compaction_interval_sec: 60 "#, ); let config = RootConfig::load_from_path("random_path.yaml"); - assert_eq!(config.worker.my_ip, "192.0.0.1"); - assert_eq!(config.worker.num_indexing_threads, 4); - assert_eq!(config.worker.pulsar_tenant, "public"); - assert_eq!(config.worker.pulsar_namespace, "default"); - assert_eq!(config.worker.kube_namespace, "chroma"); + assert_eq!(config.query_service.my_ip, "192.0.0.1"); + assert_eq!(config.query_service.my_port, 50051); + + assert_eq!(config.compaction_service.my_ip, "192.0.0.1"); + assert_eq!(config.compaction_service.my_port, 50051); Ok(()) }); } @@ -224,8 +301,10 @@ mod tests { let _ = jail.create_file( "chroma_config.yaml", r#" - worker: - num_indexing_threads: 4 + query_service: + assignment_policy: + RendezvousHashing: + hasher: Murmur3 "#, ); let _ = RootConfig::load(); @@ -239,37 +318,67 @@ mod tests { let _ = jail.create_file( "chroma_config.yaml", r#" - worker: + query_service: my_ip: "192.0.0.1" my_port: 50051 - pulsar_tenant: "public" - pulsar_namespace: "default" - kube_namespace: "chroma" - pulsar_url: "pulsar://localhost:6650" assignment_policy: RendezvousHashing: hasher: Murmur3 memberlist_provider: CustomResource: - memberlist_name: "worker-memberlist" + kube_namespace: "chroma" + memberlist_name: "query-service-memberlist" queue_size: 100 - ingest: - queue_size: 100 sysdb: Grpc: host: "localhost" port: 50051 - segment_manager: - storage_path: "/tmp" storage: S3: bucket: "chroma" - + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + + compaction_service: + my_ip: "192.0.0.1" + my_port: 50051 + assignment_policy: + RendezvousHashing: + hasher: Murmur3 + memberlist_provider: + CustomResource: + kube_namespace: "chroma" + memberlist_name: "query-service-memberlist" + queue_size: 100 + sysdb: + Grpc: + host: "localhost" + port: 50051 + storage: + S3: + bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + compactor: + compaction_manager_queue_size: 1000 + max_concurrent_jobs: 100 + compaction_interval_sec: 60 "#, ); let config = RootConfig::load(); - assert_eq!(config.worker.my_ip, "192.0.0.1"); - assert_eq!(config.worker.num_indexing_threads, num_cpus::get() as u32); + assert_eq!(config.query_service.my_ip, "192.0.0.1"); Ok(()) }); } @@ -277,44 +386,78 @@ mod tests { #[test] fn test_config_with_env_override() { Jail::expect_with(|jail| { - let _ = jail.set_env("CHROMA_WORKER__MY_IP", "192.0.0.1"); - let _ = jail.set_env("CHROMA_WORKER__MY_PORT", 50051); - let _ = jail.set_env("CHROMA_WORKER__PULSAR_TENANT", "A"); - let _ = jail.set_env("CHROMA_WORKER__PULSAR_NAMESPACE", "B"); - let _ = jail.set_env("CHROMA_WORKER__KUBE_NAMESPACE", "C"); - let _ = jail.set_env("CHROMA_WORKER__PULSAR_URL", "pulsar://localhost:6650"); + let _ = jail.set_env("CHROMA_QUERY_SERVICE__MY_IP", "192.0.0.1"); + let _ = jail.set_env("CHROMA_QUERY_SERVICE__MY_PORT", 50051); + let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__MY_IP", "192.0.0.1"); + let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__MY_PORT", 50051); let _ = jail.create_file( "chroma_config.yaml", r#" - worker: + query_service: + assignment_policy: + RendezvousHashing: + hasher: Murmur3 + memberlist_provider: + CustomResource: + kube_namespace: "chroma" + memberlist_name: "query-service-memberlist" + queue_size: 100 + sysdb: + Grpc: + host: "localhost" + port: 50051 + storage: + S3: + bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + + compaction_service: assignment_policy: RendezvousHashing: hasher: Murmur3 memberlist_provider: CustomResource: - memberlist_name: "worker-memberlist" + kube_namespace: "chroma" + memberlist_name: "compaction-service-memberlist" queue_size: 100 - ingest: - queue_size: 100 sysdb: Grpc: host: "localhost" port: 50051 - segment_manager: - storage_path: "/tmp" storage: S3: bucket: "chroma" + log: + Grpc: + host: "localhost" + port: 50051 + dispatcher: + num_worker_threads: 4 + dispatcher_queue_size: 100 + worker_queue_size: 100 + compactor: + compaction_manager_queue_size: 1000 + max_concurrent_jobs: 100 + compaction_interval_sec: 60 "#, ); let config = RootConfig::load(); - assert_eq!(config.worker.my_ip, "192.0.0.1"); - assert_eq!(config.worker.my_port, 50051); - assert_eq!(config.worker.num_indexing_threads, num_cpus::get() as u32); - assert_eq!(config.worker.pulsar_tenant, "A"); - assert_eq!(config.worker.pulsar_namespace, "B"); - assert_eq!(config.worker.kube_namespace, "C"); + assert_eq!(config.query_service.my_ip, "192.0.0.1"); + assert_eq!(config.query_service.my_port, 50051); Ok(()) }); } + + #[test] + fn test_default_config_path() { + // Sanity check that root config loads from default path correctly + let _ = RootConfig::load(); + } } diff --git a/rust/worker/src/distance/mod.rs b/rust/worker/src/distance/mod.rs new file mode 100644 index 00000000000..f8baffd363b --- /dev/null +++ b/rust/worker/src/distance/mod.rs @@ -0,0 +1,3 @@ +mod types; + +pub(crate) use types::*; diff --git a/rust/worker/src/distance/types.rs b/rust/worker/src/distance/types.rs new file mode 100644 index 00000000000..99e0df285c5 --- /dev/null +++ b/rust/worker/src/distance/types.rs @@ -0,0 +1,145 @@ +use crate::errors::{ChromaError, ErrorCodes}; +use thiserror::Error; + +/// The distance function enum. +/// # Description +/// This enum defines the distance functions supported by indices in Chroma. +/// # Variants +/// - `Euclidean` - The Euclidean or l2 norm. +/// - `Cosine` - The cosine distance. Specifically, 1 - cosine. +/// - `InnerProduct` - The inner product. Specifically, 1 - inner product. +/// # Notes +/// See https://docs.trychroma.com/usage-guide#changing-the-distance-function +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum DistanceFunction { + Euclidean, + Cosine, + InnerProduct, +} + +impl DistanceFunction { + // TOOD: Should we error if mismatched dimensions? + pub(crate) fn distance(&self, a: &[f32], b: &[f32]) -> f32 { + // TODO: implement this in SSE/AVX SIMD + // For now we write these as loops since we suspect that will more likely + // lead to the compiler vectorizing the code. (We saw this on + // Apple Silicon Macs who didn't have hand-rolled SIMD instructions in our + // C++ code). + match self { + DistanceFunction::Euclidean => { + let mut sum = 0.0; + for i in 0..a.len() { + sum += (a[i] - b[i]).powi(2); + } + sum + } + DistanceFunction::Cosine => { + // For cosine we just assume the vectors have been normalized, since that + // is what our indices expect. + let mut sum = 0.0; + for i in 0..a.len() { + sum += a[i] * b[i]; + } + 1.0_f32 - sum + } + DistanceFunction::InnerProduct => { + let mut sum = 0.0; + for i in 0..a.len() { + sum += a[i] * b[i]; + } + 1.0_f32 - sum + } + } + } +} + +#[derive(Error, Debug)] +pub(crate) enum DistanceFunctionError { + #[error("Invalid distance function `{0}`")] + InvalidDistanceFunction(String), +} + +impl ChromaError for DistanceFunctionError { + fn code(&self) -> ErrorCodes { + match self { + DistanceFunctionError::InvalidDistanceFunction(_) => ErrorCodes::InvalidArgument, + } + } +} + +impl TryFrom<&str> for DistanceFunction { + type Error = DistanceFunctionError; + + fn try_from(value: &str) -> Result { + match value { + "l2" => Ok(DistanceFunction::Euclidean), + "cosine" => Ok(DistanceFunction::Cosine), + "ip" => Ok(DistanceFunction::InnerProduct), + _ => Err(DistanceFunctionError::InvalidDistanceFunction( + value.to_string(), + )), + } + } +} + +impl Into for DistanceFunction { + fn into(self) -> String { + match self { + DistanceFunction::Euclidean => "l2".to_string(), + DistanceFunction::Cosine => "cosine".to_string(), + DistanceFunction::InnerProduct => "ip".to_string(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::convert::TryInto; + + #[test] + fn test_distance_function_try_from() { + let distance_function: DistanceFunction = "l2".try_into().unwrap(); + assert_eq!(distance_function, DistanceFunction::Euclidean); + let distance_function: DistanceFunction = "cosine".try_into().unwrap(); + assert_eq!(distance_function, DistanceFunction::Cosine); + let distance_function: DistanceFunction = "ip".try_into().unwrap(); + assert_eq!(distance_function, DistanceFunction::InnerProduct); + } + + #[test] + fn test_distance_function_into() { + let distance_function: String = DistanceFunction::Euclidean.into(); + assert_eq!(distance_function, "l2"); + let distance_function: String = DistanceFunction::Cosine.into(); + assert_eq!(distance_function, "cosine"); + let distance_function: String = DistanceFunction::InnerProduct.into(); + assert_eq!(distance_function, "ip"); + } + + #[test] + fn test_distance_function_l2sqr() { + let a = vec![1.0, 2.0, 3.0]; + let a_mag = (1.0_f32.powi(2) + 2.0_f32.powi(2) + 3.0_f32.powi(2)).sqrt(); + let a_norm = vec![1.0 / a_mag, 2.0 / a_mag, 3.0 / a_mag]; + let b = vec![4.0, 5.0, 6.0]; + let b_mag = (4.0_f32.powi(2) + 5.0_f32.powi(2) + 6.0_f32.powi(2)).sqrt(); + let b_norm = vec![4.0 / b_mag, 5.0 / b_mag, 6.0 / b_mag]; + + let l2_sqr = (1.0 - 4.0_f32).powi(2) + (2.0 - 5.0_f32).powi(2) + (3.0 - 6.0_f32).powi(2); + let inner_product_sim = 1.0_f32 + - a_norm + .iter() + .zip(b_norm.iter()) + .map(|(a, b)| a * b) + .sum::(); + + let distance_function: DistanceFunction = "l2".try_into().unwrap(); + assert_eq!(distance_function.distance(&a, &b), l2_sqr); + let distance_function: DistanceFunction = "ip".try_into().unwrap(); + assert_eq!( + distance_function.distance(&a_norm, &b_norm), + inner_product_sim + ); + } +} diff --git a/rust/worker/src/errors.rs b/rust/worker/src/errors.rs index c28d39ba9b7..086b938f265 100644 --- a/rust/worker/src/errors.rs +++ b/rust/worker/src/errors.rs @@ -1,9 +1,9 @@ // Defines 17 standard error codes based on the error codes defined in the // gRPC spec. https://grpc.github.io/grpc/core/md_doc_statuscodes.html // Custom errors can use these codes in order to allow for generic handling - use std::error::Error; +#[derive(PartialEq, Debug)] pub(crate) enum ErrorCodes { // OK is returned on success, we use "Success" since Ok is a keyword in Rust. Success = 0, @@ -41,6 +41,6 @@ pub(crate) enum ErrorCodes { DataLoss = 15, } -pub(crate) trait ChromaError: Error { +pub(crate) trait ChromaError: Error + Send { fn code(&self) -> ErrorCodes; } diff --git a/rust/worker/src/execution/config.rs b/rust/worker/src/execution/config.rs new file mode 100644 index 00000000000..d8550dc41bc --- /dev/null +++ b/rust/worker/src/execution/config.rs @@ -0,0 +1,8 @@ +use serde::Deserialize; + +#[derive(Deserialize)] +pub(crate) struct DispatcherConfig { + pub(crate) num_worker_threads: usize, + pub(crate) dispatcher_queue_size: usize, + pub(crate) worker_queue_size: usize, +} diff --git a/rust/worker/src/execution/data/data_chunk.rs b/rust/worker/src/execution/data/data_chunk.rs new file mode 100644 index 00000000000..33bd9022eaf --- /dev/null +++ b/rust/worker/src/execution/data/data_chunk.rs @@ -0,0 +1,168 @@ +use crate::types::LogRecord; +use std::sync::Arc; + +#[derive(Clone, Debug)] +pub(crate) struct DataChunk { + data: Arc<[LogRecord]>, + visibility: Arc<[bool]>, +} + +impl DataChunk { + pub fn new(data: Arc<[LogRecord]>) -> Self { + let len = data.len(); + DataChunk { + data, + visibility: vec![true; len].into(), + } + } + + /// Returns the total length of the data chunk + pub fn total_len(&self) -> usize { + self.data.len() + } + + /// Returns the number of visible elements in the data chunk + pub fn len(&self) -> usize { + self.visibility.iter().filter(|&v| *v).count() + } + + /// Returns the element at the given index + /// if the index is out of bounds, it returns None + /// # Arguments + /// * `index` - The index of the element + pub fn get(&self, index: usize) -> Option<&LogRecord> { + if index < self.data.len() { + Some(&self.data[index]) + } else { + None + } + } + + /// Returns the visibility of the element at the given index + /// if the index is out of bounds, it returns None + /// # Arguments + /// * `index` - The index of the element + pub fn get_visibility(&self, index: usize) -> Option { + if index < self.visibility.len() { + Some(self.visibility[index]) + } else { + None + } + } + + /// Sets the visibility of the elements in the data chunk. + /// Note that the length of the visibility vector should be + /// equal to the length of the data chunk. + /// + /// Note that this is the only way to change the visibility of the elements in the data chunk, + /// the data chunk does not provide a way to change the visibility of individual elements. + /// This is to ensure that the visibility of the elements is always in sync with the data. + /// If you want to change the visibility of individual elements, you should create a new data chunk. + /// + /// # Arguments + /// * `visibility` - A vector of boolean values indicating the visibility of the elements + pub fn set_visibility(&mut self, visibility: Vec) { + self.visibility = visibility.into(); + } + + /// Returns an iterator over the visible elements in the data chunk + /// The iterator returns a tuple of the element and its index + /// # Returns + /// An iterator over the visible elements in the data chunk + pub fn iter(&self) -> DataChunkIteraror<'_> { + DataChunkIteraror { + chunk: self, + index: 0, + } + } +} + +pub(crate) struct DataChunkIteraror<'a> { + chunk: &'a DataChunk, + index: usize, +} + +impl<'a> Iterator for DataChunkIteraror<'a> { + type Item = (&'a LogRecord, usize); + + fn next(&mut self) -> Option { + while self.index < self.chunk.total_len() { + let index = self.index; + match self.chunk.get_visibility(index) { + Some(true) => { + self.index += 1; + return self.chunk.get(index).map(|record| (record, index)); + } + Some(false) => { + self.index += 1; + } + None => { + break; + } + } + } + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::LogRecord; + use crate::types::Operation; + use crate::types::OperationRecord; + + #[test] + fn test_data_chunk() { + let data = vec![ + LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + ]; + let data = data.into(); + let mut chunk = DataChunk::new(data); + assert_eq!(chunk.len(), 2); + let mut iter = chunk.iter(); + let elem = iter.next(); + assert_eq!(elem.is_some(), true); + let (record, index) = elem.unwrap(); + assert_eq!(record.record.id, "embedding_id_1"); + assert_eq!(index, 0); + let elem = iter.next(); + assert_eq!(elem.is_some(), true); + let (record, index) = elem.unwrap(); + assert_eq!(record.record.id, "embedding_id_2"); + assert_eq!(index, 1); + let elem = iter.next(); + assert_eq!(elem.is_none(), true); + + let visibility = vec![true, false].into(); + chunk.set_visibility(visibility); + assert_eq!(chunk.len(), 1); + let mut iter = chunk.iter(); + let elem = iter.next(); + assert_eq!(elem.is_some(), true); + let (record, index) = elem.unwrap(); + assert_eq!(record.record.id, "embedding_id_1"); + assert_eq!(index, 0); + let elem = iter.next(); + assert_eq!(elem.is_none(), true); + } +} diff --git a/rust/worker/src/execution/data/mod.rs b/rust/worker/src/execution/data/mod.rs new file mode 100644 index 00000000000..ecbe39f3445 --- /dev/null +++ b/rust/worker/src/execution/data/mod.rs @@ -0,0 +1 @@ +pub(crate) mod data_chunk; diff --git a/rust/worker/src/execution/dispatcher.rs b/rust/worker/src/execution/dispatcher.rs new file mode 100644 index 00000000000..06007a13e57 --- /dev/null +++ b/rust/worker/src/execution/dispatcher.rs @@ -0,0 +1,290 @@ +use super::{operator::TaskMessage, worker_thread::WorkerThread}; +use crate::execution::config::DispatcherConfig; +use crate::{ + config::Configurable, + errors::ChromaError, + system::{Component, ComponentContext, Handler, Receiver, System}, +}; +use async_trait::async_trait; +use std::fmt::Debug; + +/// The dispatcher is responsible for distributing tasks to worker threads. +/// It is a component that receives tasks and distributes them to worker threads. +/** +```plaintext + ┌─────────────────────────────────────────┐ + │ │ + │ │ + │ │ + TaskMessage ───────────►├─────┐ Dispatcher │ + │ ▼ │ + │ ┌┬───────────────────────────────┐ │ + │ │┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼┼│ │ + └────┴──────────────┴─────────────────┴───┘ + ▲ + │ │ + │ │ + TaskRequestMessage │ │ TaskMessage + │ │ + │ │ + ▼ + ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ + │ │ │ │ │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + │ Worker │ │ Worker │ │ Worker │ + │ │ │ │ │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + └────────────────┘ └────────────────┘ └────────────────┘ +``` +## Implementation notes +- The dispatcher has a queue of tasks that it distributes to worker threads +- A worker thread sends a TaskRequestMessage to the dispatcher when it is ready for a new task +- If no task is available for the worker thread, the dispatcher will place that worker's reciever + in a queue and send a task to the worker when it recieves another one +- The reason to introduce this abstraction is to allow us to control fairness and dynamically adjust + system utilization. It also makes mechanisms like pausing/stopping work easier. + It would have likely been more performant to use the Tokio MT runtime, but we chose to use + this abstraction to grant us flexibility. We can always switch to Tokio MT later if we need to, + or make this dispatcher much more performant through implementing memory-awareness, task-batches, + coarser work-stealing, and other optimizations. +*/ +#[derive(Debug)] +pub(crate) struct Dispatcher { + task_queue: Vec, + waiters: Vec, + n_worker_threads: usize, + queue_size: usize, + worker_queue_size: usize, +} + +impl Dispatcher { + /// Create a new dispatcher + /// # Parameters + /// - n_worker_threads: The number of worker threads to use + /// - queue_size: The size of the components message queue + /// - worker_queue_size: The size of the worker components queue + pub fn new(n_worker_threads: usize, queue_size: usize, worker_queue_size: usize) -> Self { + Dispatcher { + task_queue: Vec::new(), + waiters: Vec::new(), + n_worker_threads, + queue_size, + worker_queue_size, + } + } + + /// Spawn worker threads + /// # Parameters + /// - system: The system to spawn the worker threads in + /// - self_receiver: The receiver to send tasks to the worker threads, this is a address back to the dispatcher + fn spawn_workers( + &self, + system: &mut System, + self_receiver: Box>, + ) { + for _ in 0..self.n_worker_threads { + let worker = WorkerThread::new(self_receiver.clone(), self.worker_queue_size); + system.start_component(worker); + } + } + + /// Enqueue a task to be processed + /// # Parameters + /// - task: The task to enqueue + async fn enqueue_task(&mut self, task: TaskMessage) { + // If a worker is waiting for a task, send it to the worker in FIFO order + // Otherwise, add it to the task queue + match self.waiters.pop() { + Some(channel) => match channel.reply_to.send(task).await { + Ok(_) => {} + Err(e) => { + println!("Error sending task to worker: {:?}", e); + } + }, + None => { + self.task_queue.push(task); + } + } + } + + /// Handle a work request from a worker thread + /// # Parameters + /// - worker: The request for work + /// If no work is available, the worker will be placed in a queue and a task will be sent to it + /// when one is available + async fn handle_work_request(&mut self, request: TaskRequestMessage) { + match self.task_queue.pop() { + Some(task) => match request.reply_to.send(task).await { + Ok(_) => {} + Err(e) => { + println!("Error sending task to worker: {:?}", e); + } + }, + None => { + self.waiters.push(request); + } + } + } +} + +#[async_trait] +impl Configurable for Dispatcher { + async fn try_from_config(config: &DispatcherConfig) -> Result> { + Ok(Dispatcher::new( + config.num_worker_threads, + config.dispatcher_queue_size, + config.worker_queue_size, + )) + } +} + +/// A message that a worker thread sends to the dispatcher to request a task +/// # Members +/// - reply_to: The receiver to send the task to, this is the worker thread +#[derive(Debug)] +pub(super) struct TaskRequestMessage { + reply_to: Box>, +} + +impl TaskRequestMessage { + /// Create a new TaskRequestMessage + /// # Parameters + /// - reply_to: The receiver to send the task to, this is the worker thread + /// that is requesting the task + pub(super) fn new(reply_to: Box>) -> Self { + TaskRequestMessage { reply_to } + } +} + +// ============= Component implementation ============= + +#[async_trait] +impl Component for Dispatcher { + fn queue_size(&self) -> usize { + self.queue_size + } + + async fn on_start(&mut self, ctx: &ComponentContext) { + self.spawn_workers(&mut ctx.system.clone(), ctx.sender.as_receiver()); + } +} + +#[async_trait] +impl Handler for Dispatcher { + async fn handle(&mut self, task: TaskMessage, _ctx: &ComponentContext) { + self.enqueue_task(task).await; + } +} + +// Worker sends a request for task +#[async_trait] +impl Handler for Dispatcher { + async fn handle(&mut self, message: TaskRequestMessage, _ctx: &ComponentContext) { + self.handle_work_request(message).await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + execution::operator::{wrap, Operator}, + system::System, + }; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + + // Create a component that will schedule DISPATCH_COUNT invocations of the MockOperator + // on an interval of DISPATCH_FREQUENCY_MS. + // Each invocation will sleep for MOCK_OPERATOR_SLEEP_DURATION_MS to simulate work + // Use THREAD_COUNT worker threads + const MOCK_OPERATOR_SLEEP_DURATION_MS: u64 = 100; + const DISPATCH_FREQUENCY_MS: u64 = 5; + const DISPATCH_COUNT: usize = 50; + const THREAD_COUNT: usize = 4; + + #[derive(Debug)] + struct MockOperator {} + #[async_trait] + impl Operator for MockOperator { + type Error = (); + async fn run(&self, input: &f32) -> Result { + // sleep to simulate work + tokio::time::sleep(tokio::time::Duration::from_millis( + MOCK_OPERATOR_SLEEP_DURATION_MS, + )) + .await; + Ok(input.to_string()) + } + } + + #[derive(Debug)] + struct MockDispatchUser { + pub dispatcher: Box>, + counter: Arc, // We expect to recieve DISPATCH_COUNT messages + } + #[async_trait] + impl Component for MockDispatchUser { + fn queue_size(&self) -> usize { + 1000 + } + + async fn on_start(&mut self, ctx: &ComponentContext) { + // dispatch a new task every DISPATCH_FREQUENCY_MS for DISPATCH_COUNT times + let duration = std::time::Duration::from_millis(DISPATCH_FREQUENCY_MS); + ctx.scheduler.schedule_interval( + ctx.sender.clone(), + (), + duration, + Some(DISPATCH_COUNT), + ctx, + ); + } + } + #[async_trait] + impl Handler> for MockDispatchUser { + async fn handle( + &mut self, + _message: Result, + ctx: &ComponentContext, + ) { + self.counter.fetch_add(1, Ordering::SeqCst); + let curr_count = self.counter.load(Ordering::SeqCst); + // Cancel self + if curr_count == DISPATCH_COUNT { + ctx.cancellation_token.cancel(); + } + } + } + + #[async_trait] + impl Handler<()> for MockDispatchUser { + async fn handle(&mut self, _message: (), ctx: &ComponentContext) { + let task = wrap(Box::new(MockOperator {}), 42.0, ctx.sender.as_receiver()); + let res = self.dispatcher.send(task).await; + } + } + + #[tokio::test] + async fn test_dispatcher() { + let system = System::new(); + let dispatcher = Dispatcher::new(THREAD_COUNT, 1000, 1000); + let dispatcher_handle = system.start_component(dispatcher); + let counter = Arc::new(AtomicUsize::new(0)); + let dispatch_user = MockDispatchUser { + dispatcher: dispatcher_handle.receiver(), + counter: counter.clone(), + }; + let mut dispatch_user_handle = system.start_component(dispatch_user); + // yield to allow the component to process the messages + tokio::task::yield_now().await; + // Join on the dispatch user, since it will kill itself after DISPATCH_COUNT messages + dispatch_user_handle.join().await; + // We should have received DISPATCH_COUNT messages + assert_eq!(counter.load(Ordering::SeqCst), DISPATCH_COUNT); + } +} diff --git a/rust/worker/src/execution/mod.rs b/rust/worker/src/execution/mod.rs new file mode 100644 index 00000000000..1d361780d77 --- /dev/null +++ b/rust/worker/src/execution/mod.rs @@ -0,0 +1,7 @@ +pub(crate) mod config; +mod data; +pub(crate) mod dispatcher; +pub(crate) mod operator; +mod operators; +pub(crate) mod orchestration; +mod worker_thread; diff --git a/rust/worker/src/execution/operator.rs b/rust/worker/src/execution/operator.rs new file mode 100644 index 00000000000..935c01eb16e --- /dev/null +++ b/rust/worker/src/execution/operator.rs @@ -0,0 +1,75 @@ +use crate::system::Receiver; +use async_trait::async_trait; +use std::fmt::Debug; + +/// An operator takes a generic input and returns a generic output. +/// It is a definition of a function. +#[async_trait] +pub(super) trait Operator: Send + Sync + Debug +where + I: Send + Sync, + O: Send + Sync, +{ + type Error; + // It would have been nice to do this with a default trait for result + // but that's not stable in rust yet. + async fn run(&self, input: &I) -> Result; +} + +/// A task is a wrapper around an operator and its input. +/// It is a description of a function to be run. +#[derive(Debug)] +struct Task +where + Input: Send + Sync + Debug, + Output: Send + Sync + Debug, +{ + operator: Box>, + input: Input, + reply_channel: Box>>, +} + +/// A message type used by the dispatcher to send tasks to worker threads. +pub(crate) type TaskMessage = Box; + +/// A task wrapper is a trait that can be used to run a task. We use it to +/// erase the I, O types from the Task struct so that tasks. +#[async_trait] +pub(crate) trait TaskWrapper: Send + Debug { + async fn run(&self); +} + +/// Implement the TaskWrapper trait for every Task. This allows us to +/// erase the I, O types from the Task struct so that tasks can be +/// stored in a homogenous queue regardless of their input and output types. +#[async_trait] +impl TaskWrapper for Task +where + Error: Debug, + Input: Send + Sync + Debug, + Output: Send + Sync + Debug, +{ + async fn run(&self) { + let output = self.operator.run(&self.input).await; + let res = self.reply_channel.send(output).await; + // TODO: if this errors, it means the caller was dropped + } +} + +/// Wrap an operator and its input into a task message. +pub(super) fn wrap( + operator: Box>, + input: Input, + reply_channel: Box>>, +) -> TaskMessage +where + Error: Debug + 'static, + Input: Send + Sync + Debug + 'static, + Output: Send + Sync + Debug + 'static, +{ + Box::new(Task { + operator, + input, + reply_channel, + }) +} diff --git a/rust/worker/src/execution/operators/brute_force_knn.rs b/rust/worker/src/execution/operators/brute_force_knn.rs new file mode 100644 index 00000000000..2ebaa9f79a5 --- /dev/null +++ b/rust/worker/src/execution/operators/brute_force_knn.rs @@ -0,0 +1,275 @@ +use crate::execution::data::data_chunk::DataChunk; +use crate::{distance::DistanceFunction, execution::operator::Operator}; +use async_trait::async_trait; +use std::cmp::Ord; +use std::cmp::Ordering; +use std::cmp::PartialOrd; +use std::collections::BinaryHeap; + +/// The brute force k-nearest neighbors operator is responsible for computing the k-nearest neighbors +/// of a given query vector against a set of vectors using brute force calculation. +/// # Note +/// - Callers should ensure that the input vectors are normalized if using the cosine similarity metric. +#[derive(Debug)] +pub struct BruteForceKnnOperator {} + +/// The input to the brute force k-nearest neighbors operator. +/// # Parameters +/// * `data` - The vectors to query against. +/// * `query` - The query vector. +/// * `k` - The number of nearest neighbors to find. +/// * `distance_metric` - The distance metric to use. +#[derive(Debug)] +pub struct BruteForceKnnOperatorInput { + pub data: DataChunk, + pub query: Vec, + pub k: usize, + pub distance_metric: DistanceFunction, +} + +/// The output of the brute force k-nearest neighbors operator. +/// # Parameters +/// * `data` - The vectors to query against. Only the vectors that are nearest neighbors are visible. +/// * `indices` - The indices of the nearest neighbors. This is a mask against the `query_vecs` input. +/// One row for each query vector. +/// * `distances` - The distances of the nearest neighbors. +/// One row for each query vector. +#[derive(Debug)] +pub struct BruteForceKnnOperatorOutput { + pub data: DataChunk, + pub indices: Vec, + pub distances: Vec, +} + +pub type BruteForceKnnOperatorResult = Result; + +#[derive(Debug)] +struct Entry { + index: usize, + distance: f32, +} + +impl Ord for Entry { + fn cmp(&self, other: &Self) -> Ordering { + if self.distance == other.distance { + Ordering::Equal + } else if self.distance > other.distance { + // This is a min heap, so we need to reverse the ordering. + Ordering::Less + } else { + // This is a min heap, so we need to reverse the ordering. + Ordering::Greater + } + } +} + +impl PartialOrd for Entry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl PartialEq for Entry { + fn eq(&self, other: &Self) -> bool { + self.distance == other.distance + } +} + +impl Eq for Entry {} + +#[async_trait] +impl Operator for BruteForceKnnOperator { + type Error = (); + + async fn run(&self, input: &BruteForceKnnOperatorInput) -> BruteForceKnnOperatorResult { + let mut heap = BinaryHeap::with_capacity(input.k); + let data_chunk = &input.data; + for data in data_chunk.iter() { + let log_record = data.0; + let index = data.1; + + let embedding = match &log_record.record.embedding { + Some(embedding) => embedding, + None => { + continue; + } + }; + let distance = input.distance_metric.distance(&embedding[..], &input.query); + heap.push(Entry { index, distance }); + } + + let mut visibility = vec![false; data_chunk.total_len()]; + let mut sorted_indices = Vec::with_capacity(input.k); + let mut sorted_distances = Vec::with_capacity(input.k); + let mut i = 0; + while i < input.k { + let entry = match heap.pop() { + Some(entry) => entry, + None => { + break; + } + }; + sorted_indices.push(entry.index); + sorted_distances.push(entry.distance); + visibility[entry.index] = true; + i += 1; + } + let mut data_chunk = data_chunk.clone(); + data_chunk.set_visibility(visibility); + Ok(BruteForceKnnOperatorOutput { + data: data_chunk, + indices: sorted_indices, + distances: sorted_distances, + }) + } +} + +#[cfg(test)] +mod tests { + use crate::types::LogRecord; + use crate::types::Operation; + use crate::types::OperationRecord; + + use super::*; + + #[tokio::test] + async fn test_brute_force_knn_l2sqr() { + let operator = BruteForceKnnOperator {}; + let data = vec![ + LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: Some(vec![0.0, 0.0, 0.0]), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: Some(vec![0.0, 1.0, 1.0]), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 3, + record: OperationRecord { + id: "embedding_id_3".to_string(), + embedding: Some(vec![7.0, 8.0, 9.0]), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + ]; + let data_chunk = DataChunk::new(data.into()); + + let input = BruteForceKnnOperatorInput { + data: data_chunk, + query: vec![0.0, 0.0, 0.0], + k: 2, + distance_metric: DistanceFunction::Euclidean, + }; + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.indices, vec![0, 1]); + let distance_1 = 0.0_f32.powi(2) + 1.0_f32.powi(2) + 1.0_f32.powi(2); + assert_eq!(output.distances, vec![0.0, distance_1]); + assert_eq!(output.data.get_visibility(0), Some(true)); + assert_eq!(output.data.get_visibility(1), Some(true)); + assert_eq!(output.data.get_visibility(2), Some(false)); + } + + #[tokio::test] + async fn test_brute_force_knn_cosine() { + let operator = BruteForceKnnOperator {}; + + let norm_1 = (1.0_f32.powi(2) + 2.0_f32.powi(2) + 3.0_f32.powi(2)).sqrt(); + let data_1 = vec![1.0 / norm_1, 2.0 / norm_1, 3.0 / norm_1]; + + let norm_2 = (0.0_f32.powi(2) + -1.0_f32.powi(2) + 6.0_f32.powi(2)).sqrt(); + let data_2 = vec![0.0 / norm_2, -1.0 / norm_2, 6.0 / norm_2]; + let data = vec![ + LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: Some(vec![0.0, 1.0, 0.0]), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: Some(data_1.clone()), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 3, + record: OperationRecord { + id: "embedding_id_3".to_string(), + embedding: Some(data_2.clone()), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + ]; + let data_chunk = DataChunk::new(data.into()); + + let input = BruteForceKnnOperatorInput { + data: data_chunk, + query: vec![0.0, 1.0, 0.0], + k: 2, + distance_metric: DistanceFunction::InnerProduct, + }; + let output = operator.run(&input).await.unwrap(); + + assert_eq!(output.indices, vec![0, 1]); + let expected_distance_1 = + 1.0f32 - ((data_1[0] * 0.0) + (data_1[1] * 1.0) + (data_1[2] * 0.0)); + assert_eq!(output.distances, vec![0.0, expected_distance_1]); + assert_eq!(output.data.get_visibility(0), Some(true)); + assert_eq!(output.data.get_visibility(1), Some(true)); + assert_eq!(output.data.get_visibility(2), Some(false)); + } + + #[tokio::test] + async fn test_data_less_than_k() { + // If we have less data than k, we should return all the data, sorted by distance. + let operator = BruteForceKnnOperator {}; + let data = vec![LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: Some(vec![0.0, 0.0, 0.0]), + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }]; + + let data_chunk = DataChunk::new(data.into()); + + let input = BruteForceKnnOperatorInput { + data: data_chunk, + query: vec![0.0, 0.0, 0.0], + k: 2, + distance_metric: DistanceFunction::Euclidean, + }; + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.indices, vec![0]); + assert_eq!(output.distances, vec![0.0]); + assert_eq!(output.data.get_visibility(0), Some(true)); + } +} diff --git a/rust/worker/src/execution/operators/mod.rs b/rust/worker/src/execution/operators/mod.rs new file mode 100644 index 00000000000..ed31dca33db --- /dev/null +++ b/rust/worker/src/execution/operators/mod.rs @@ -0,0 +1,4 @@ +pub(super) mod brute_force_knn; +pub(super) mod normalize_vectors; +pub(super) mod partition; +pub(super) mod pull_log; diff --git a/rust/worker/src/execution/operators/normalize_vectors.rs b/rust/worker/src/execution/operators/normalize_vectors.rs new file mode 100644 index 00000000000..57607ddd2ee --- /dev/null +++ b/rust/worker/src/execution/operators/normalize_vectors.rs @@ -0,0 +1,81 @@ +use crate::execution::operator::Operator; +use async_trait::async_trait; + +const EPS: f32 = 1e-30; + +#[derive(Debug)] +pub struct NormalizeVectorOperator {} + +pub struct NormalizeVectorOperatorInput { + pub vectors: Vec>, +} + +pub struct NormalizeVectorOperatorOutput { + pub normalized_vectors: Vec>, +} + +#[async_trait] +impl Operator + for NormalizeVectorOperator +{ + type Error = (); + + async fn run( + &self, + input: &NormalizeVectorOperatorInput, + ) -> Result { + // TODO: this should not have to reallocate the vectors. We can optimize this later. + let mut normalized_vectors = Vec::with_capacity(input.vectors.len()); + for vector in &input.vectors { + let mut norm = 0.0; + for x in vector { + norm += x * x; + } + let norm = 1.0 / (norm.sqrt() + EPS); + let normalized_vector = vector.iter().map(|x| x * norm).collect(); + normalized_vectors.push(normalized_vector); + } + Ok(NormalizeVectorOperatorOutput { normalized_vectors }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const COMPARE_EPS: f32 = 1e-9; + fn float_eps_eq(a: &[f32], b: &[f32]) -> bool { + a.iter() + .zip(b.iter()) + .all(|(a, b)| (a - b).abs() < COMPARE_EPS) + } + + #[tokio::test] + async fn test_normalize_vector() { + let operator = NormalizeVectorOperator {}; + let input = NormalizeVectorOperatorInput { + vectors: vec![ + vec![1.0, 2.0, 3.0], + vec![4.0, 5.0, 6.0], + vec![7.0, 8.0, 9.0], + ], + }; + + let output = operator.run(&input).await.unwrap(); + let expected_output = NormalizeVectorOperatorOutput { + normalized_vectors: vec![ + vec![0.26726124, 0.5345225, 0.8017837], + vec![0.45584232, 0.5698029, 0.6837635], + vec![0.5025707, 0.5743665, 0.64616233], + ], + }; + + for (a, b) in output + .normalized_vectors + .iter() + .zip(expected_output.normalized_vectors.iter()) + { + assert!(float_eps_eq(a, b), "{:?} != {:?}", a, b); + } + } +} diff --git a/rust/worker/src/execution/operators/partition.rs b/rust/worker/src/execution/operators/partition.rs new file mode 100644 index 00000000000..fec691993f0 --- /dev/null +++ b/rust/worker/src/execution/operators/partition.rs @@ -0,0 +1,218 @@ +use crate::errors::{ChromaError, ErrorCodes}; +use crate::execution::data::data_chunk::DataChunk; +use crate::execution::operator::Operator; +use async_trait::async_trait; +use std::collections::HashMap; +use thiserror::Error; + +#[derive(Debug)] +/// The partition Operator takes a DataChunk and presents a copy-free +/// view of N partitions by breaking the data into partitions by max_partition_size. It will group operations +/// on the same key into the same partition. Due to this, the max_partition_size is a +/// soft-limit, since if there are more operations to a key than max_partition_size we cannot +/// partition the data. +pub struct PartitionOperator {} + +/// The input to the partition operator. +/// # Parameters +/// * `records` - The records to partition. +#[derive(Debug)] +pub struct PartitionInput { + pub(crate) records: DataChunk, + pub(crate) max_partition_size: usize, +} + +impl PartitionInput { + /// Create a new partition input. + /// # Parameters + /// * `records` - The records to partition. + /// * `max_partition_size` - The maximum size of a partition. Since we are trying to + /// partition the records by id, which can casue the partition size to be larger than this + /// value. + pub fn new(records: DataChunk, max_partition_size: usize) -> Self { + PartitionInput { + records, + max_partition_size, + } + } +} + +/// The output of the partition operator. +/// # Parameters +/// * `records` - The partitioned records. +#[derive(Debug)] +pub struct PartitionOutput { + pub(crate) records: Vec, +} + +#[derive(Debug, Error)] +pub enum PartitionError { + #[error("Failed to partition records.")] + PartitionError, +} + +impl ChromaError for PartitionError { + fn code(&self) -> ErrorCodes { + match self { + PartitionError::PartitionError => ErrorCodes::Internal, + } + } +} + +pub type PartitionResult = Result; + +impl PartitionOperator { + pub fn new() -> Box { + Box::new(PartitionOperator {}) + } + + pub fn partition(&self, records: &DataChunk, partition_size: usize) -> Vec { + let mut map = HashMap::new(); + for data in records.iter() { + let log_record = data.0; + let index = data.1; + let key = log_record.record.id.clone(); + map.entry(key).or_insert_with(Vec::new).push(index); + } + let mut result = Vec::new(); + // Create a new DataChunk for each parition of records with partition_size without + // data copying. + let mut current_batch_size = 0; + let mut new_partition = true; + let mut visibility = vec![false; records.total_len()]; + for (_, v) in map.iter() { + // create DataChunk with partition_size by masking the visibility of the records + // in the partition. + if new_partition { + visibility = vec![false; records.total_len()]; + new_partition = false; + } + for i in v.iter() { + visibility[*i] = true; + } + current_batch_size += v.len(); + if current_batch_size >= partition_size { + let mut new_data_chunk = records.clone(); + new_data_chunk.set_visibility(visibility.clone()); + result.push(new_data_chunk); + new_partition = true; + current_batch_size = 0; + } + } + // handle the case that the last group is smaller than the group_size. + if !new_partition { + let mut new_data_chunk = records.clone(); + new_data_chunk.set_visibility(visibility.clone()); + result.push(new_data_chunk); + } + result + } + + fn determine_partition_size(&self, num_records: usize, threshold: usize) -> usize { + if num_records < threshold { + return num_records; + } else { + return threshold; + } + } +} + +#[async_trait] +impl Operator for PartitionOperator { + type Error = PartitionError; + + async fn run(&self, input: &PartitionInput) -> PartitionResult { + let records = &input.records; + let partition_size = self.determine_partition_size(records.len(), input.max_partition_size); + let deduped_records = self.partition(records, partition_size); + return Ok(PartitionOutput { + records: deduped_records, + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::{LogRecord, Operation, OperationRecord}; + use std::sync::Arc; + + #[tokio::test] + async fn test_partition_operator() { + let data = vec![ + LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + LogRecord { + log_offset: 3, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + ]; + let data: Arc<[LogRecord]> = data.into(); + + // Test group size is larger than the number of records + let chunk = DataChunk::new(data.clone()); + let operator = PartitionOperator::new(); + let input = PartitionInput::new(chunk, 4); + let result = operator.run(&input).await.unwrap(); + assert_eq!(result.records.len(), 1); + assert_eq!(result.records[0].len(), 3); + + // Test group size is the same as the number of records + let chunk = DataChunk::new(data.clone()); + let operator = PartitionOperator::new(); + let input = PartitionInput::new(chunk, 3); + let result = operator.run(&input).await.unwrap(); + assert_eq!(result.records.len(), 1); + assert_eq!(result.records[0].len(), 3); + + // Test group size is smaller than the number of records + let chunk = DataChunk::new(data.clone()); + let operator = PartitionOperator::new(); + let input = PartitionInput::new(chunk, 2); + let mut result = operator.run(&input).await.unwrap(); + + // The result can be 1 or 2 groups depending on the order of the records. + assert!(result.records.len() == 2 || result.records.len() == 1); + if result.records.len() == 2 { + result.records.sort_by(|a, b| a.len().cmp(&b.len())); + assert_eq!(result.records[0].len(), 1); + assert_eq!(result.records[1].len(), 2); + } else { + assert_eq!(result.records[0].len(), 3); + } + + // Test group size is smaller than the number of records + let chunk = DataChunk::new(data.clone()); + let operator = PartitionOperator::new(); + let input = PartitionInput::new(chunk, 1); + let mut result = operator.run(&input).await.unwrap(); + assert_eq!(result.records.len(), 2); + result.records.sort_by(|a, b| a.len().cmp(&b.len())); + assert_eq!(result.records[0].len(), 1); + assert_eq!(result.records[1].len(), 2); + } +} diff --git a/rust/worker/src/execution/operators/pull_log.rs b/rust/worker/src/execution/operators/pull_log.rs new file mode 100644 index 00000000000..f7b5486fbb4 --- /dev/null +++ b/rust/worker/src/execution/operators/pull_log.rs @@ -0,0 +1,248 @@ +use crate::execution::data::data_chunk::DataChunk; +use crate::execution::operator::Operator; +use crate::log::log::Log; +use crate::log::log::PullLogsError; +use async_trait::async_trait; +use uuid::Uuid; + +/// The pull logs operator is responsible for reading logs from the log service. +#[derive(Debug)] +pub struct PullLogsOperator { + client: Box, +} + +impl PullLogsOperator { + /// Create a new pull logs operator. + /// # Parameters + /// * `client` - The log client to use for reading logs. + pub fn new(client: Box) -> Box { + Box::new(PullLogsOperator { client }) + } +} + +/// The input to the pull logs operator. +/// # Parameters +/// * `collection_id` - The collection id to read logs from. +/// * `offset` - The offset to start reading logs from. +/// * `batch_size` - The number of log entries to read. +/// * `num_records` - The maximum number of records to read. +/// * `end_timestamp` - The end timestamp to read logs until. +#[derive(Debug)] +pub struct PullLogsInput { + collection_id: Uuid, + offset: i64, + batch_size: i32, + num_records: Option, + end_timestamp: Option, +} + +impl PullLogsInput { + /// Create a new pull logs input. + /// # Parameters + /// * `collection_id` - The collection id to read logs from. + /// * `offset` - The offset to start reading logs from. + /// * `batch_size` - The number of log entries to read. + /// * `num_records` - The maximum number of records to read. + /// * `end_timestamp` - The end timestamp to read logs until. + pub fn new( + collection_id: Uuid, + offset: i64, + batch_size: i32, + num_records: Option, + end_timestamp: Option, + ) -> Self { + PullLogsInput { + collection_id, + offset, + batch_size, + num_records, + end_timestamp, + } + } +} + +/// The output of the pull logs operator. +#[derive(Debug)] +pub struct PullLogsOutput { + logs: DataChunk, +} + +impl PullLogsOutput { + /// Create a new pull logs output. + /// # Parameters + /// * `logs` - The logs that were read. + pub fn new(logs: DataChunk) -> Self { + PullLogsOutput { logs } + } + + /// Get the log entries that were read by an invocation of the pull logs operator. + /// # Returns + /// The log entries that were read. + pub fn logs(&self) -> DataChunk { + self.logs.clone() + } +} + +pub type PullLogsResult = Result; + +#[async_trait] +impl Operator for PullLogsOperator { + type Error = PullLogsError; + + async fn run(&self, input: &PullLogsInput) -> PullLogsResult { + // We expect the log to be cheaply cloneable, we need to clone it since we need + // a mutable reference to it. Not necessarily the best, but it works for our needs. + let mut client_clone = self.client.clone(); + let batch_size = input.batch_size; + let mut num_records_read = 0; + let mut offset = input.offset; + let mut result = Vec::new(); + loop { + let logs = client_clone + .read(input.collection_id, offset, batch_size, input.end_timestamp) + .await; + + let mut logs = match logs { + Ok(logs) => logs, + Err(e) => { + return Err(e); + } + }; + + if logs.is_empty() { + break; + } + + num_records_read += logs.len(); + offset += batch_size as i64; + result.append(&mut logs); + + // We used a a timestamp and we didn't get a full batch, so we have retrieved + // the last batch of logs relevant to our query + if input.end_timestamp.is_some() && num_records_read < batch_size as usize { + break; + } + + // We have read all the records up to the size we wanted + if input.num_records.is_some() + && num_records_read >= input.num_records.unwrap() as usize + { + break; + } + } + if input.num_records.is_some() && result.len() > input.num_records.unwrap() as usize { + result.truncate(input.num_records.unwrap() as usize); + } + // Convert to DataChunk + let data_chunk = DataChunk::new(result.into()); + Ok(PullLogsOutput::new(data_chunk)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::log::log::InMemoryLog; + use crate::log::log::InternalLogRecord; + use crate::types::LogRecord; + use crate::types::Operation; + use crate::types::OperationRecord; + use std::str::FromStr; + use uuid::Uuid; + + #[tokio::test] + async fn test_pull_logs() { + let mut log = Box::new(InMemoryLog::new()); + + let collection_uuid_1 = Uuid::from_str("00000000-0000-0000-0000-000000000001").unwrap(); + let collection_id_1 = collection_uuid_1.to_string(); + log.add_log( + collection_id_1.clone(), + Box::new(InternalLogRecord { + collection_id: collection_id_1.clone(), + log_offset: 1, + log_ts: 1, + record: LogRecord { + log_offset: 1, + record: OperationRecord { + id: "embedding_id_1".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + }), + ); + log.add_log( + collection_id_1.clone(), + Box::new(InternalLogRecord { + collection_id: collection_id_1.clone(), + log_offset: 2, + log_ts: 2, + record: LogRecord { + log_offset: 2, + record: OperationRecord { + id: "embedding_id_2".to_string(), + embedding: None, + encoding: None, + metadata: None, + operation: Operation::Add, + }, + }, + }), + ); + + let operator = PullLogsOperator::new(log); + + // Pull all logs from collection 1 + let input = PullLogsInput::new(collection_uuid_1, 0, 1, None, None); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 2); + + // Pull all logs from collection 1 with a large batch size + let input = PullLogsInput::new(collection_uuid_1, 0, 100, None, None); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 2); + + // Pull logs from collection 1 with a limit + let input = PullLogsInput::new(collection_uuid_1, 0, 1, Some(1), None); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 1); + + // Pull logs from collection 1 with an end timestamp + let input = PullLogsInput::new(collection_uuid_1, 0, 1, None, Some(1)); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 1); + + // Pull logs from collection 1 with an end timestamp + let input = PullLogsInput::new(collection_uuid_1, 0, 1, None, Some(2)); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 2); + + // Pull logs from collection 1 with an end timestamp and a limit + let input = PullLogsInput::new(collection_uuid_1, 0, 1, Some(1), Some(2)); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 1); + + // Pull logs from collection 1 with a limit and a large batch size + let input = PullLogsInput::new(collection_uuid_1, 0, 100, Some(1), None); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 1); + + // Pull logs from collection 1 with an end timestamp and a large batch size + let input = PullLogsInput::new(collection_uuid_1, 0, 100, None, Some(1)); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 1); + + // Pull logs from collection 1 with an end timestamp and a large batch size + let input = PullLogsInput::new(collection_uuid_1, 0, 100, None, Some(2)); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 2); + + // Pull logs from collection 1 with an end timestamp and a limit and a large batch size + let input = PullLogsInput::new(collection_uuid_1, 0, 100, Some(1), Some(2)); + let output = operator.run(&input).await.unwrap(); + assert_eq!(output.logs().len(), 1); + } +} diff --git a/rust/worker/src/execution/orchestration/compact.rs b/rust/worker/src/execution/orchestration/compact.rs new file mode 100644 index 00000000000..85fe0802249 --- /dev/null +++ b/rust/worker/src/execution/orchestration/compact.rs @@ -0,0 +1,232 @@ +use super::super::operator::{wrap, TaskMessage}; +use crate::compactor::CompactionJob; +use crate::errors::ChromaError; +use crate::execution::data::data_chunk::DataChunk; +use crate::execution::operators::partition::PartitionInput; +use crate::execution::operators::partition::PartitionOperator; +use crate::execution::operators::partition::PartitionResult; +use crate::execution::operators::pull_log::PullLogsInput; +use crate::execution::operators::pull_log::PullLogsOperator; +use crate::execution::operators::pull_log::PullLogsResult; +use crate::log::log::Log; +use crate::system::Component; +use crate::system::Handler; +use crate::system::Receiver; +use crate::system::System; +use async_trait::async_trait; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; +use uuid::Uuid; + +/** The state of the orchestrator. +In chroma, we have a relatively fixed number of query plans that we can execute. Rather +than a flexible state machine abstraction, we just manually define the states that we +expect to encounter for a given query plan. This is a bit more rigid, but it's also simpler and easier to +understand. We can always add more abstraction later if we need it. +```plaintext + + ┌───► Write─────-------┐ + │ │ + Pending ─► PullLogs ─► Group │ ├─► Flush ─► Finished + │ │ + └───► Write ───────────┘ + +``` +*/ +#[derive(Debug)] +enum ExecutionState { + Pending, + PullLogs, + Partition, + Write, + Flush, + Finished, +} + +#[derive(Debug)] +pub struct CompactOrchestrator { + id: Uuid, + task: CompactionJob, + state: ExecutionState, + // Component Execution + system: System, + collection_id: Uuid, + // Dependencies + log: Box, + // Dispatcher + dispatcher: Box>, + // Result Channel + result_channel: + Option>>>, +} + +// TODO: we need to improve this response +#[derive(Debug)] +pub struct CompactionResponse { + id: Uuid, + task: CompactionJob, + message: String, +} + +impl CompactOrchestrator { + pub fn new( + task: CompactionJob, + system: System, + collection_id: Uuid, + log: Box, + dispatcher: Box>, + result_channel: Option< + tokio::sync::oneshot::Sender>>, + >, + ) -> Self { + CompactOrchestrator { + id: Uuid::new_v4(), + task, + state: ExecutionState::Pending, + system, + collection_id, + log, + dispatcher, + result_channel, + } + } + + async fn pull_logs(&mut self, self_address: Box>) { + self.state = ExecutionState::PullLogs; + let operator = PullLogsOperator::new(self.log.clone()); + let collection_id = self.collection_id; + let end_timestamp = SystemTime::now().duration_since(UNIX_EPOCH); + let end_timestamp = match end_timestamp { + // TODO: change protobuf definition to use u64 instead of i64 + Ok(end_timestamp) => end_timestamp.as_secs() as i64, + Err(e) => { + // Log an error and reply + return + return; + } + }; + let input = PullLogsInput::new(collection_id, 0, 100, None, Some(end_timestamp)); + let task = wrap(operator, input, self_address); + match self.dispatcher.send(task).await { + Ok(_) => (), + Err(e) => { + // TODO: log an error and reply to caller + } + } + } + + async fn partition( + &mut self, + records: DataChunk, + self_address: Box>, + ) { + self.state = ExecutionState::Partition; + // TODO: make this configurable + let max_partition_size = 100; + let operator = PartitionOperator::new(); + let input = PartitionInput::new(records, max_partition_size); + let task = wrap(operator, input, self_address); + match self.dispatcher.send(task).await { + Ok(_) => (), + Err(e) => { + // TODO: log an error and reply to caller + } + } + } + + async fn write(&mut self, records: Vec) { + self.state = ExecutionState::Write; + + for record in records { + // TODO: implement write + } + } + + pub(crate) async fn run(mut self) -> Result> { + let (tx, rx) = tokio::sync::oneshot::channel(); + self.result_channel = Some(tx); + let mut handle = self.system.clone().start_component(self); + let result = rx.await; + handle.stop(); + result.unwrap() + } +} + +// ============== Component Implementation ============== + +#[async_trait] +impl Component for CompactOrchestrator { + fn queue_size(&self) -> usize { + 1000 // TODO: make configurable + } + + async fn on_start(&mut self, ctx: &crate::system::ComponentContext) -> () { + self.pull_logs(ctx.sender.as_receiver()).await; + } +} + +// ============== Handlers ============== +#[async_trait] +impl Handler for CompactOrchestrator { + async fn handle( + &mut self, + message: PullLogsResult, + ctx: &crate::system::ComponentContext, + ) { + let records = match message { + Ok(result) => result.logs(), + Err(e) => { + // Log an error and return + let result_channel = match self.result_channel.take() { + Some(tx) => tx, + None => { + // Log an error + return; + } + }; + let _ = result_channel.send(Err(Box::new(e))); + return; + } + }; + self.partition(records, ctx.sender.as_receiver()).await; + } +} + +#[async_trait] +impl Handler for CompactOrchestrator { + async fn handle( + &mut self, + message: PartitionResult, + _ctx: &crate::system::ComponentContext, + ) { + let records = match message { + Ok(result) => result.records, + Err(e) => { + // Log an error and return + let result_channel = match self.result_channel.take() { + Some(tx) => tx, + None => { + // Log an error + return; + } + }; + let _ = result_channel.send(Err(Box::new(e))); + return; + } + }; + // TODO: implement write records + // For now, we will return to execution state to the compaction manager + let result_channel = match self.result_channel.take() { + Some(tx) => tx, + None => { + // Log an error + return; + } + }; + let response = CompactionResponse { + id: self.id, + task: self.task.clone(), + message: "Compaction Complete".to_string(), + }; + let _ = result_channel.send(Ok(response)); + } +} diff --git a/rust/worker/src/execution/orchestration/hnsw.rs b/rust/worker/src/execution/orchestration/hnsw.rs new file mode 100644 index 00000000000..5cce87d6be4 --- /dev/null +++ b/rust/worker/src/execution/orchestration/hnsw.rs @@ -0,0 +1,246 @@ +use super::super::operator::{wrap, TaskMessage}; +use super::super::operators::pull_log::{PullLogsInput, PullLogsOperator}; +use crate::distance::DistanceFunction; +use crate::errors::ChromaError; +use crate::execution::operators::brute_force_knn::{ + BruteForceKnnOperator, BruteForceKnnOperatorInput, BruteForceKnnOperatorResult, +}; +use crate::execution::operators::pull_log::PullLogsResult; +use crate::sysdb::sysdb::SysDb; +use crate::system::System; +use crate::types::VectorQueryResult; +use crate::{ + log::log::Log, + system::{Component, Handler, Receiver}, +}; +use async_trait::async_trait; +use std::fmt::Debug; +use std::time::{SystemTime, UNIX_EPOCH}; +use uuid::Uuid; + +/** The state of the orchestrator. +In chroma, we have a relatively fixed number of query plans that we can execute. Rather +than a flexible state machine abstraction, we just manually define the states that we +expect to encounter for a given query plan. This is a bit more rigid, but it's also simpler and easier to +understand. We can always add more abstraction later if we need it. +```plaintext + + ┌───► Brute Force ─────┐ + │ │ + Pending ─► PullLogs ─► Group│ ├─► MergeResults ─► Finished + │ │ + └───► HNSW ────────────┘ + +``` +*/ +#[derive(Debug)] +enum ExecutionState { + Pending, + PullLogs, + Partition, + QueryKnn, + MergeResults, + Finished, +} + +#[derive(Debug)] +pub(crate) struct HnswQueryOrchestrator { + state: ExecutionState, + // Component Execution + system: System, + // Query state + query_vectors: Vec>, + k: i32, + include_embeddings: bool, + segment_id: Uuid, + // Services + log: Box, + sysdb: Box, + dispatcher: Box>, + // Result channel + result_channel: Option< + tokio::sync::oneshot::Sender>, Box>>, + >, +} + +impl HnswQueryOrchestrator { + pub(crate) fn new( + system: System, + query_vectors: Vec>, + k: i32, + include_embeddings: bool, + segment_id: Uuid, + log: Box, + sysdb: Box, + dispatcher: Box>, + ) -> Self { + HnswQueryOrchestrator { + state: ExecutionState::Pending, + system, + query_vectors, + k, + include_embeddings, + segment_id, + log, + sysdb, + dispatcher, + result_channel: None, + } + } + + /// Get the collection id for a segment id. + /// TODO: This can be cached + async fn get_collection_id_for_segment_id(&mut self, segment_id: Uuid) -> Option { + let segments = self + .sysdb + .get_segments(Some(segment_id), None, None, None) + .await; + match segments { + Ok(segments) => match segments.get(0) { + Some(segment) => segment.collection, + None => None, + }, + Err(e) => { + // Log an error and return + return None; + } + } + } + + async fn pull_logs(&mut self, self_address: Box>) { + self.state = ExecutionState::PullLogs; + let operator = PullLogsOperator::new(self.log.clone()); + let collection_id = match self.get_collection_id_for_segment_id(self.segment_id).await { + Some(collection_id) => collection_id, + None => { + // Log an error and reply + return + return; + } + }; + let end_timestamp = SystemTime::now().duration_since(UNIX_EPOCH); + let end_timestamp = match end_timestamp { + // TODO: change protobuf definition to use u64 instead of i64 + Ok(end_timestamp) => end_timestamp.as_nanos() as i64, + Err(e) => { + // Log an error and reply + return + return; + } + }; + let input = PullLogsInput::new(collection_id, 0, 100, None, Some(end_timestamp)); + let task = wrap(operator, input, self_address); + match self.dispatcher.send(task).await { + Ok(_) => (), + Err(e) => { + // TODO: log an error and reply to caller + } + } + } + + /// Run the orchestrator and return the result. + /// # Note + /// Use this over spawning the component directly. This method will start the component and + /// wait for it to finish before returning the result. + pub(crate) async fn run(mut self) -> Result>, Box> { + let (tx, rx) = tokio::sync::oneshot::channel(); + self.result_channel = Some(tx); + let mut handle = self.system.clone().start_component(self); + let result = rx.await; + handle.stop(); + result.unwrap() + } +} + +// ============== Component Implementation ============== + +#[async_trait] +impl Component for HnswQueryOrchestrator { + fn queue_size(&self) -> usize { + 1000 // TODO: make configurable + } + + async fn on_start(&mut self, ctx: &crate::system::ComponentContext) -> () { + self.pull_logs(ctx.sender.as_receiver()).await; + } +} + +// ============== Handlers ============== + +#[async_trait] +impl Handler for HnswQueryOrchestrator { + async fn handle( + &mut self, + message: PullLogsResult, + ctx: &crate::system::ComponentContext, + ) { + self.state = ExecutionState::Partition; + + // TODO: implement the remaining state transitions and operators + // TODO: don't need all this cloning and data shuffling, once we land the chunk abstraction + match message { + Ok(logs) => { + let bf_input = BruteForceKnnOperatorInput { + data: logs.logs(), + query: self.query_vectors[0].clone(), + k: self.k as usize, + distance_metric: DistanceFunction::Euclidean, + }; + let operator = Box::new(BruteForceKnnOperator {}); + let task = wrap(operator, bf_input, ctx.sender.as_receiver()); + match self.dispatcher.send(task).await { + Ok(_) => (), + Err(e) => { + // TODO: log an error and reply to caller + } + } + } + Err(e) => { + // Log an error + return; + } + } + } +} + +#[async_trait] +impl Handler for HnswQueryOrchestrator { + async fn handle( + &mut self, + message: BruteForceKnnOperatorResult, + _ctx: &crate::system::ComponentContext, + ) { + // This is an example of the final state transition and result + let result_channel = match self.result_channel.take() { + Some(tx) => tx, + None => { + // Log an error + return; + } + }; + + match message { + Ok(output) => { + let mut result = Vec::new(); + let mut query_results = Vec::new(); + for (index, distance) in output.indices.iter().zip(output.distances.iter()) { + let query_result = VectorQueryResult { + id: index.to_string(), + distance: *distance, + vector: None, + }; + query_results.push(query_result); + } + result.push(query_results); + + match result_channel.send(Ok(result)) { + Ok(_) => (), + Err(e) => { + // Log an error + } + } + } + Err(_) => { + // Log an error + } + } + } +} diff --git a/rust/worker/src/execution/orchestration/mod.rs b/rust/worker/src/execution/orchestration/mod.rs new file mode 100644 index 00000000000..d91d03f3e75 --- /dev/null +++ b/rust/worker/src/execution/orchestration/mod.rs @@ -0,0 +1,4 @@ +mod compact; +mod hnsw; +pub(crate) use compact::*; +pub(crate) use hnsw::*; diff --git a/rust/worker/src/execution/worker_thread.rs b/rust/worker/src/execution/worker_thread.rs new file mode 100644 index 00000000000..d651a725d34 --- /dev/null +++ b/rust/worker/src/execution/worker_thread.rs @@ -0,0 +1,58 @@ +use super::{dispatcher::TaskRequestMessage, operator::TaskMessage}; +use crate::system::{Component, ComponentContext, ComponentRuntime, Handler, Receiver}; +use async_trait::async_trait; +use std::fmt::{Debug, Formatter, Result}; + +/// A worker thread is responsible for executing tasks +/// It sends requests to the dispatcher for new tasks. +/// # Implementation notes +/// - The actor loop will block until work is available +pub(super) struct WorkerThread { + dispatcher: Box>, + queue_size: usize, +} + +impl WorkerThread { + pub(super) fn new( + dispatcher: Box>, + queue_size: usize, + ) -> WorkerThread { + WorkerThread { + dispatcher, + queue_size, + } + } +} + +impl Debug for WorkerThread { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + f.debug_struct("WorkerThread").finish() + } +} + +#[async_trait] +impl Component for WorkerThread { + fn queue_size(&self) -> usize { + self.queue_size + } + + fn runtime() -> ComponentRuntime { + ComponentRuntime::Dedicated + } + + async fn on_start(&mut self, ctx: &ComponentContext) -> () { + let req = TaskRequestMessage::new(ctx.sender.as_receiver()); + let res = self.dispatcher.send(req).await; + // TODO: what to do with resp? + } +} + +#[async_trait] +impl Handler for WorkerThread { + async fn handle(&mut self, task: TaskMessage, ctx: &ComponentContext) { + task.run().await; + let req: TaskRequestMessage = TaskRequestMessage::new(ctx.sender.as_receiver()); + let res = self.dispatcher.send(req).await; + // TODO: task run should be able to error and we should send it as part of the result + } +} diff --git a/rust/worker/src/index/fulltext/mod.rs b/rust/worker/src/index/fulltext/mod.rs new file mode 100644 index 00000000000..ec2bac7d473 --- /dev/null +++ b/rust/worker/src/index/fulltext/mod.rs @@ -0,0 +1,2 @@ +pub mod tokenizer; +mod types; diff --git a/rust/worker/src/index/fulltext/tokenizer.rs b/rust/worker/src/index/fulltext/tokenizer.rs new file mode 100644 index 00000000000..9109d27af21 --- /dev/null +++ b/rust/worker/src/index/fulltext/tokenizer.rs @@ -0,0 +1,82 @@ +use tantivy::tokenizer::{NgramTokenizer, Token, TokenStream, Tokenizer}; + +pub(crate) trait ChromaTokenStream { + fn process(&mut self, sink: &mut dyn FnMut(&Token)); + fn get_tokens(&self) -> &Vec; +} + +pub(crate) struct TantivyChromaTokenStream { + tokens: Vec, +} + +impl TantivyChromaTokenStream { + pub fn new(tokens: Vec) -> Self { + TantivyChromaTokenStream { tokens } + } +} + +impl ChromaTokenStream for TantivyChromaTokenStream { + fn process(&mut self, sink: &mut dyn FnMut(&Token)) { + for token in &self.tokens { + sink(token); + } + } + + fn get_tokens(&self) -> &Vec { + &self.tokens + } +} + +pub(crate) trait ChromaTokenizer { + fn encode(&mut self, text: &str) -> Box; +} + +pub(crate) struct TantivyChromaTokenizer { + tokenizer: Box, +} + +impl TantivyChromaTokenizer { + pub fn new(tokenizer: Box) -> Self { + TantivyChromaTokenizer { tokenizer } + } +} + +impl ChromaTokenizer for TantivyChromaTokenizer { + fn encode(&mut self, text: &str) -> Box { + let mut token_stream = self.tokenizer.token_stream(text); + let mut tokens = Vec::new(); + token_stream.process(&mut |token| { + tokens.push(token.clone()); + }); + Box::new(TantivyChromaTokenStream::new(tokens)) + } +} + +mod test { + use super::*; + + #[test] + fn test_chroma_tokenizer() { + let tokenizer: Box = Box::new(NgramTokenizer::new(1, 1, false).unwrap()); + let mut chroma_tokenizer = TantivyChromaTokenizer::new(tokenizer); + let mut token_stream = chroma_tokenizer.encode("hello world"); + let mut tokens = Vec::new(); + token_stream.process(&mut |token| { + tokens.push(token.clone()); + }); + assert_eq!(tokens.len(), 11); + assert_eq!(tokens[0].text, "h"); + assert_eq!(tokens[1].text, "e"); + } + + #[test] + fn test_get_tokens() { + let tokenizer: Box = Box::new(NgramTokenizer::new(1, 1, false).unwrap()); + let mut chroma_tokenizer = TantivyChromaTokenizer::new(tokenizer); + let token_stream = chroma_tokenizer.encode("hello world"); + let tokens = token_stream.get_tokens(); + assert_eq!(tokens.len(), 11); + assert_eq!(tokens[0].text, "h"); + assert_eq!(tokens[1].text, "e"); + } +} diff --git a/rust/worker/src/index/fulltext/types.rs b/rust/worker/src/index/fulltext/types.rs new file mode 100644 index 00000000000..49ea5138d11 --- /dev/null +++ b/rust/worker/src/index/fulltext/types.rs @@ -0,0 +1,342 @@ +use crate::blockstore::{Blockfile, BlockfileKey, Key, PositionalPostingListBuilder, Value}; +use crate::errors::{ChromaError, ErrorCodes}; +use crate::index::fulltext::tokenizer::ChromaTokenizer; +use std::collections::HashMap; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum FullTextIndexError { + #[error("Already in a transaction")] + AlreadyInTransaction, + #[error("Not in a transaction")] + NotInTransaction, +} + +impl ChromaError for FullTextIndexError { + fn code(&self) -> ErrorCodes { + match self { + FullTextIndexError::AlreadyInTransaction => ErrorCodes::FailedPrecondition, + FullTextIndexError::NotInTransaction => ErrorCodes::FailedPrecondition, + } + } +} + +pub(crate) trait FullTextIndex { + fn begin_transaction(&mut self) -> Result<(), Box>; + fn commit_transaction(&mut self) -> Result<(), Box>; + + // Must be done inside a transaction. + fn add_document(&mut self, document: &str, offset_id: i32) -> Result<(), Box>; + // Only searches committed state. + fn search(&mut self, query: &str) -> Result, Box>; +} + +pub(crate) struct BlockfileFullTextIndex { + posting_lists_blockfile: Box, + frequencies_blockfile: Box, + tokenizer: Box, + in_transaction: bool, + + // term -> positional posting list builder for that term + uncommitted: HashMap, + uncommitted_frequencies: HashMap, +} + +impl BlockfileFullTextIndex { + pub(crate) fn new( + posting_lists_blockfile: Box, + frequencies_blockfile: Box, + tokenizer: Box, + ) -> Self { + BlockfileFullTextIndex { + posting_lists_blockfile, + frequencies_blockfile, + tokenizer, + in_transaction: false, + uncommitted: HashMap::new(), + uncommitted_frequencies: HashMap::new(), + } + } +} + +impl FullTextIndex for BlockfileFullTextIndex { + fn begin_transaction(&mut self) -> Result<(), Box> { + if self.in_transaction { + return Err(Box::new(FullTextIndexError::AlreadyInTransaction)); + } + self.posting_lists_blockfile.begin_transaction()?; + self.frequencies_blockfile.begin_transaction()?; + self.in_transaction = true; + Ok(()) + } + + fn commit_transaction(&mut self) -> Result<(), Box> { + if !self.in_transaction { + return Err(Box::new(FullTextIndexError::NotInTransaction)); + } + self.in_transaction = false; + for (key, mut value) in self.uncommitted.drain() { + let positional_posting_list = value.build(); + let blockfilekey = BlockfileKey::new("".to_string(), Key::String(key.to_string())); + self.posting_lists_blockfile.set( + blockfilekey, + Value::PositionalPostingListValue(positional_posting_list), + ); + } + for (key, value) in self.uncommitted_frequencies.drain() { + let blockfilekey = BlockfileKey::new("".to_string(), Key::String(key.to_string())); + self.frequencies_blockfile + .set(blockfilekey, Value::IntValue(value)); + } + self.posting_lists_blockfile.commit_transaction()?; + self.frequencies_blockfile.commit_transaction()?; + self.uncommitted.clear(); + Ok(()) + } + + fn add_document(&mut self, document: &str, offset_id: i32) -> Result<(), Box> { + if !self.in_transaction { + return Err(Box::new(FullTextIndexError::NotInTransaction)); + } + let tokens = self.tokenizer.encode(document); + for token in tokens.get_tokens() { + self.uncommitted_frequencies + .entry(token.text.to_string()) + .and_modify(|e| *e += 1) + .or_insert(1); + let builder = self + .uncommitted + .entry(token.text.to_string()) + .or_insert(PositionalPostingListBuilder::new()); + + // Store starting positions of tokens. These are NOT affected by token filters. + // For search, we can use the start and end positions to compute offsets to + // check full string match. + // + // See https://docs.rs/tantivy/latest/tantivy/tokenizer/struct.Token.html + if !builder.contains_doc_id(offset_id) { + // Casting to i32 is safe since we limit the size of the document. + builder.add_doc_id_and_positions(offset_id, vec![token.offset_from as i32]); + } else { + builder.add_positions_for_doc_id(offset_id, vec![token.offset_from as i32]); + } + } + Ok(()) + } + + fn search(&mut self, query: &str) -> Result, Box> { + let binding = self.tokenizer.encode(query); + let tokens = binding.get_tokens(); + + // Get query tokens sorted by frequency. + let mut token_frequencies = vec![]; + for token in tokens { + let blockfilekey = + BlockfileKey::new("".to_string(), Key::String(token.text.to_string())); + let value = self.frequencies_blockfile.get(blockfilekey); + match value { + Ok(Value::IntValue(frequency)) => { + token_frequencies.push((token.text.to_string(), frequency)); + } + Ok(_) => { + return Ok(vec![]); + } + Err(_) => { + // TODO error handling from blockfile + return Ok(vec![]); + } + } + } + token_frequencies.sort_by(|a, b| a.1.cmp(&b.1)); + + // Populate initial candidates with the least-frequent token's posting list. + // doc ID -> possible starting locations for the query. + let mut candidates: HashMap> = HashMap::new(); + let blockfilekey = + BlockfileKey::new("".to_string(), Key::String(tokens[0].text.to_string())); + let first_token_positional_posting_list = + match self.posting_lists_blockfile.get(blockfilekey).unwrap() { + Value::PositionalPostingListValue(arr) => arr, + _ => panic!("Value is not an arrow struct array"), + }; + let first_token_offset = tokens[0].offset_from as i32; + for doc_id in first_token_positional_posting_list.get_doc_ids().values() { + let positions = first_token_positional_posting_list + .get_positions_for_doc_id(*doc_id) + .unwrap(); + let positions_vec: Vec = positions + .values() + .iter() + .map(|x| *x - first_token_offset) + .collect(); + candidates.insert(*doc_id, positions_vec); + } + + // Iterate through the rest of the tokens, intersecting the posting lists with the candidates. + for (token, _) in token_frequencies[1..].iter() { + let blockfilekey = BlockfileKey::new("".to_string(), Key::String(token.to_string())); + let positional_posting_list = + match self.posting_lists_blockfile.get(blockfilekey).unwrap() { + Value::PositionalPostingListValue(arr) => arr, + _ => panic!("Value is not an arrow struct array"), + }; + let token_offset = tokens + .iter() + .find(|t| t.text == *token) + .unwrap() + .offset_from as i32; + let mut new_candidates: HashMap> = HashMap::new(); + for (doc_id, positions) in candidates.iter() { + let mut new_positions = vec![]; + for position in positions { + if let Some(positions_for_doc_id) = + positional_posting_list.get_positions_for_doc_id(*doc_id) + { + for position_for_doc_id in positions_for_doc_id.values() { + if position_for_doc_id - token_offset == *position { + new_positions.push(*position); + } + } + } + } + if !new_positions.is_empty() { + new_candidates.insert(*doc_id, new_positions); + } + } + candidates = new_candidates; + } + + let mut results = vec![]; + for (doc_id, _) in candidates.drain() { + results.push(doc_id); + } + + Ok(results) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::blockstore::provider::{BlockfileProvider, HashMapBlockfileProvider}; + use crate::blockstore::{KeyType, ValueType}; + use crate::index::fulltext::tokenizer::TantivyChromaTokenizer; + use tantivy::tokenizer::NgramTokenizer; + + #[test] + fn test_new() { + let mut provider = HashMapBlockfileProvider::new(); + let pl_blockfile = provider + .create("pl", KeyType::String, ValueType::PositionalPostingList) + .unwrap(); + let freq_blockfile = provider + .create("freq", KeyType::String, ValueType::Int) + .unwrap(); + let tokenizer = Box::new(TantivyChromaTokenizer::new(Box::new( + NgramTokenizer::new(1, 1, false).unwrap(), + ))); + let _index = BlockfileFullTextIndex::new(pl_blockfile, freq_blockfile, tokenizer); + } + + #[test] + fn test_index_single_document() { + let mut provider = HashMapBlockfileProvider::new(); + let pl_blockfile = provider + .create("pl", KeyType::String, ValueType::PositionalPostingList) + .unwrap(); + let freq_blockfile = provider + .create("freq", KeyType::String, ValueType::Int) + .unwrap(); + let tokenizer = Box::new(TantivyChromaTokenizer::new(Box::new( + NgramTokenizer::new(1, 1, false).unwrap(), + ))); + let mut index = BlockfileFullTextIndex::new(pl_blockfile, freq_blockfile, tokenizer); + index.begin_transaction().unwrap(); + index.add_document("hello world", 1).unwrap(); + index.commit_transaction().unwrap(); + + let res = index.search("hello"); + assert_eq!(res.unwrap(), vec![1]); + } + + #[test] + fn test_search_absent_token() { + let mut provider = HashMapBlockfileProvider::new(); + let pl_blockfile = provider + .create("pl", KeyType::String, ValueType::PositionalPostingList) + .unwrap(); + let freq_blockfile = provider + .create("freq", KeyType::String, ValueType::Int) + .unwrap(); + let tokenizer = Box::new(TantivyChromaTokenizer::new(Box::new( + NgramTokenizer::new(1, 1, false).unwrap(), + ))); + let mut index = BlockfileFullTextIndex::new(pl_blockfile, freq_blockfile, tokenizer); + index.begin_transaction().unwrap(); + index.add_document("hello world", 1).unwrap(); + index.commit_transaction().unwrap(); + + let res = index.search("chroma"); + assert!(res.unwrap().is_empty()); + } + + #[test] + fn test_index_and_search_multiple_documents() { + let mut provider = HashMapBlockfileProvider::new(); + let pl_blockfile = provider + .create("pl", KeyType::String, ValueType::PositionalPostingList) + .unwrap(); + let freq_blockfile = provider + .create("freq", KeyType::String, ValueType::Int) + .unwrap(); + let tokenizer = Box::new(TantivyChromaTokenizer::new(Box::new( + NgramTokenizer::new(1, 1, false).unwrap(), + ))); + let mut index = BlockfileFullTextIndex::new(pl_blockfile, freq_blockfile, tokenizer); + index.begin_transaction().unwrap(); + index.add_document("hello world", 1).unwrap(); + index.add_document("hello chroma", 2).unwrap(); + index.add_document("chroma world", 3).unwrap(); + index.commit_transaction().unwrap(); + + let res = index.search("hello").unwrap(); + assert!(res.contains(&1)); + assert!(res.contains(&2)); + + let res = index.search("world").unwrap(); + assert!(res.contains(&1)); + assert!(res.contains(&3)); + + let res = index.search("llo chro").unwrap(); + assert!(res.contains(&2)); + } + + #[test] + fn test_special_characters_search() { + let mut provider = HashMapBlockfileProvider::new(); + let pl_blockfile = provider + .create("pl", KeyType::String, ValueType::PositionalPostingList) + .unwrap(); + let freq_blockfile = provider + .create("freq", KeyType::String, ValueType::Int) + .unwrap(); + let tokenizer = Box::new(TantivyChromaTokenizer::new(Box::new( + NgramTokenizer::new(1, 1, false).unwrap(), + ))); + let mut index = BlockfileFullTextIndex::new(pl_blockfile, freq_blockfile, tokenizer); + index.begin_transaction().unwrap(); + index.add_document("!!!!", 1).unwrap(); + index.add_document(",,!!", 2).unwrap(); + index.add_document(".!", 3).unwrap(); + index.add_document("!.!.!.!", 4).unwrap(); + index.commit_transaction().unwrap(); + + let res = index.search("!!").unwrap(); + assert!(res.contains(&1)); + assert!(res.contains(&2)); + + let res = index.search(".!").unwrap(); + assert!(res.contains(&3)); + assert!(res.contains(&4)); + } +} diff --git a/rust/worker/src/index/hnsw.rs b/rust/worker/src/index/hnsw.rs index 49eb5efb2c9..45f72be388e 100644 --- a/rust/worker/src/index/hnsw.rs +++ b/rust/worker/src/index/hnsw.rs @@ -307,7 +307,7 @@ extern "C" { pub mod test { use super::*; - use crate::index::types::DistanceFunction; + use crate::distance::DistanceFunction; use crate::index::utils; use rand::Rng; use rayon::prelude::*; @@ -321,7 +321,7 @@ pub mod test { let tmp_dir = tempdir().unwrap(); let persist_path = tmp_dir.path().to_str().unwrap().to_string(); let distance_function = DistanceFunction::Euclidean; - let mut index = HnswIndex::init( + let index = HnswIndex::init( &IndexConfig { dimensionality: d as i32, distance_function: distance_function, @@ -382,9 +382,9 @@ pub mod test { let mut rng: rand::prelude::ThreadRng = rand::thread_rng(); let mut datas = Vec::new(); - for i in 0..n { + for _ in 0..n { let mut data: Vec = Vec::new(); - for i in 0..960 { + for _ in 0..960 { data.push(rng.gen()); } datas.push(data); diff --git a/rust/worker/src/index/metadata/mod.rs b/rust/worker/src/index/metadata/mod.rs new file mode 100644 index 00000000000..b316bb795e5 --- /dev/null +++ b/rust/worker/src/index/metadata/mod.rs @@ -0,0 +1,3 @@ +mod types; + +// TODO reexport the types module \ No newline at end of file diff --git a/rust/worker/src/index/metadata/types.rs b/rust/worker/src/index/metadata/types.rs new file mode 100644 index 00000000000..6ff0195aae0 --- /dev/null +++ b/rust/worker/src/index/metadata/types.rs @@ -0,0 +1,762 @@ +use crate::blockstore::provider::BlockfileProvider; +use crate::blockstore::{Blockfile, BlockfileKey, Key, Value}; +use crate::errors::{ChromaError, ErrorCodes}; + +use roaring::RoaringBitmap; +use std::{collections::HashMap, marker::PhantomData}; +use thiserror::Error; + +#[derive(Debug, Error)] +pub(crate) enum MetadataIndexError { + #[error("Key not found")] + NotFoundError, + #[error("This operation cannot be done in a transaction")] + InTransaction, + #[error("This operation can only be done in a transaction")] + NotInTransaction, +} + +impl ChromaError for MetadataIndexError { + fn code(&self) -> ErrorCodes { + match self { + MetadataIndexError::NotFoundError => ErrorCodes::InvalidArgument, + MetadataIndexError::InTransaction => ErrorCodes::InvalidArgument, + MetadataIndexError::NotInTransaction => ErrorCodes::InvalidArgument, + } + } +} + +pub(crate) trait MetadataIndexValue { + fn to_blockfile_key(&self) -> Key; +} +impl MetadataIndexValue for String { + fn to_blockfile_key(&self) -> Key { + Key::String(self.clone()) + } +} +impl MetadataIndexValue for f32 { + fn to_blockfile_key(&self) -> Key { + Key::Float(*self) + } +} +impl MetadataIndexValue for bool { + fn to_blockfile_key(&self) -> Key { + Key::Bool(*self) + } +} + +pub(crate) trait MetadataIndex { + fn begin_transaction(&mut self) -> Result<(), Box>; + fn commit_transaction(&mut self) -> Result<(), Box>; + fn in_transaction(&self) -> bool; + + // Must be in a transaction to put or delete. + fn set(&mut self, key: &str, value: T, offset_id: u32) -> Result<(), Box>; + // Can delete anything -- if it's not in committed state the delete will be silently discarded. + fn delete(&mut self, key: &str, value: T, offset_id: u32) -> Result<(), Box>; + + // Always reads from committed state. + fn get(&self, key: &str, value: T) -> Result>; +} + +struct BlockfileMetadataIndex { + metadata_value_type: PhantomData, + blockfile: Box, + in_transaction: bool, + uncommitted_rbms: HashMap, +} + +impl BlockfileMetadataIndex { + pub fn new(init_blockfile: Box) -> Self { + BlockfileMetadataIndex { + metadata_value_type: PhantomData, + blockfile: init_blockfile, + in_transaction: false, + uncommitted_rbms: HashMap::new(), + } + } + + fn look_up_key_and_populate_uncommitted_rbms( + &mut self, + key: &BlockfileKey, + ) -> Result<(), Box> { + if !self.uncommitted_rbms.contains_key(&key) { + match self.blockfile.get(key.clone()) { + Ok(Value::RoaringBitmapValue(rbm)) => { + self.uncommitted_rbms.insert(key.clone(), rbm); + } + _ => { + let rbm = RoaringBitmap::new(); + self.uncommitted_rbms.insert(key.clone(), rbm); + } + }; + } + Ok(()) + } +} + +impl MetadataIndex for BlockfileMetadataIndex { + fn begin_transaction(&mut self) -> Result<(), Box> { + if self.in_transaction { + return Err(Box::new(MetadataIndexError::InTransaction)); + } + self.blockfile.begin_transaction()?; + self.in_transaction = true; + Ok(()) + } + + fn commit_transaction(&mut self) -> Result<(), Box> { + if !self.in_transaction { + return Err(Box::new(MetadataIndexError::NotInTransaction)); + } + for (key, rbm) in self.uncommitted_rbms.drain() { + self.blockfile + .set(key.clone(), Value::RoaringBitmapValue(rbm.clone())); + } + self.blockfile.commit_transaction()?; + self.in_transaction = false; + self.uncommitted_rbms.clear(); + Ok(()) + } + + fn in_transaction(&self) -> bool { + self.in_transaction + } + + fn set(&mut self, key: &str, value: T, offset_id: u32) -> Result<(), Box> { + if !self.in_transaction { + return Err(Box::new(MetadataIndexError::NotInTransaction)); + } + let blockfilekey = BlockfileKey::new(key.to_string(), value.to_blockfile_key()); + self.look_up_key_and_populate_uncommitted_rbms(&blockfilekey)?; + let rbm = self.uncommitted_rbms.get_mut(&blockfilekey).unwrap(); + rbm.insert(offset_id); + Ok(()) + } + + fn delete(&mut self, key: &str, value: T, offset_id: u32) -> Result<(), Box> { + if !self.in_transaction { + return Err(Box::new(MetadataIndexError::NotInTransaction)); + } + let blockfilekey = BlockfileKey::new(key.to_string(), value.to_blockfile_key()); + self.look_up_key_and_populate_uncommitted_rbms(&blockfilekey)?; + let rbm = self.uncommitted_rbms.get_mut(&blockfilekey).unwrap(); + rbm.remove(offset_id); + Ok(()) + } + + fn get(&self, key: &str, value: T) -> Result> { + if self.in_transaction { + return Err(Box::new(MetadataIndexError::InTransaction)); + } + let blockfilekey = BlockfileKey::new(key.to_string(), value.to_blockfile_key()); + match self.blockfile.get(blockfilekey) { + Ok(Value::RoaringBitmapValue(rbm)) => Ok(rbm), + _ => Ok(RoaringBitmap::new()), + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::blockstore::provider::HashMapBlockfileProvider; + use crate::blockstore::{KeyType, ValueType}; + + #[test] + fn test_string_value_metadata_index_error_when_not_in_transaction() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::new(blockfile); + let result = index.set("key", "value".to_string(), 1); + assert_eq!(result.is_err(), true); + let result = index.delete("key", "value".to_string(), 1); + assert_eq!(result.is_err(), true); + let result = index.commit_transaction(); + assert_eq!(result.is_err(), true); + } + + #[test] + fn test_string_value_metadata_index_empty_transaction() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.commit_transaction().unwrap(); + } + + #[test] + fn test_string_value_metadata_index_set_get() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", "value".to_string(), 1).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", "value".to_string()).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(1), true); + } + + #[test] + fn test_float_value_metadata_index_set_get() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", 1.0, 1).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", 1.0).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(1), true); + } + + #[test] + fn test_bool_value_metadata_index_set_get() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", true, 1).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", true).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(1), true); + } + + #[test] + fn test_string_value_metadata_index_set_delete_get() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", "value".to_string(), 1).unwrap(); + index.delete("key", "value".to_string(), 1).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", "value".to_string()).unwrap(); + assert_eq!(bitmap.len(), 0); + } + + #[test] + fn test_string_value_metadata_index_set_delete_set_get() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", "value".to_string(), 1).unwrap(); + index.delete("key", "value".to_string(), 1).unwrap(); + index.set("key", "value".to_string(), 1).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", "value".to_string()).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(1), true); + } + + #[test] + fn test_string_value_metadata_index_multiple_keys() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key1", "value".to_string(), 1).unwrap(); + index.set("key2", "value".to_string(), 2).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key1", "value".to_string()).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(1), true); + + let bitmap = index.get("key2", "value".to_string()).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(2), true); + } + + #[test] + fn test_string_value_metadata_index_multiple_values() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", "value1".to_string(), 1).unwrap(); + index.set("key", "value2".to_string(), 2).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", "value1".to_string()).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(1), true); + + let bitmap = index.get("key", "value2".to_string()).unwrap(); + assert_eq!(bitmap.len(), 1); + assert_eq!(bitmap.contains(2), true); + } + + #[test] + fn test_string_value_metadata_index_delete_in_standalone_transaction() { + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + let mut index = BlockfileMetadataIndex::::new(blockfile); + index.begin_transaction().unwrap(); + index.set("key", "value".to_string(), 1).unwrap(); + index.commit_transaction().unwrap(); + + index.begin_transaction().unwrap(); + index.delete("key", "value".to_string(), 1).unwrap(); + index.commit_transaction().unwrap(); + + let bitmap = index.get("key", "value".to_string()).unwrap(); + assert_eq!(bitmap.len(), 0); + } + + use proptest::prelude::*; + use proptest::test_runner::Config; + use proptest_state_machine::{prop_state_machine, ReferenceStateMachine, StateMachineTest}; + use std::rc::Rc; + + // Utility function to check if a Vec and RoaringBitmap contain equivalent + // sets. + fn vec_rbm_eq(a: &Vec, b: &RoaringBitmap) -> bool { + if a.len() != b.len() as usize { + return false; + } + for offset in a { + if !b.contains(*offset) { + return false; + } + } + for offset in b { + if !a.contains(&offset) { + return false; + } + } + return true; + } + + pub(crate) trait PropTestValue: + MetadataIndexValue + PartialEq + Eq + Clone + std::hash::Hash + std::fmt::Debug + { + fn strategy() -> BoxedStrategy; + } + + impl PropTestValue for String { + fn strategy() -> BoxedStrategy { + ".{0,10}".prop_map(|x| x.to_string()).boxed() + } + } + + impl PropTestValue for bool { + fn strategy() -> BoxedStrategy { + prop_oneof![Just(true), Just(false)].boxed() + } + } + + // f32 doesn't implement Hash or Eq so we need to wrap it to use + // in our reference state machine's HashMap. This is a bit of a hack + // but only used in tests. + #[derive(Clone, Debug, PartialEq)] + struct FloatWrapper(f32); + + impl std::hash::Hash for FloatWrapper { + fn hash(&self, state: &mut H) { + self.0.to_bits().hash(state); + } + } + + impl Eq for FloatWrapper {} + + impl MetadataIndexValue for FloatWrapper { + fn to_blockfile_key(&self) -> Key { + self.0.to_blockfile_key() + } + } + + impl PropTestValue for FloatWrapper { + fn strategy() -> BoxedStrategy { + (0..1000).prop_map(|x| FloatWrapper(x as f32)).boxed() + } + } + + #[derive(Clone, Debug)] + pub(crate) enum Transition { + BeginTransaction, + CommitTransaction, + Set(String, T, u32), + Delete(String, T, u32), + Get(String, T), + } + + #[derive(Clone, Debug)] + pub(crate) struct ReferenceState { + // Are we in a transaction? + in_transaction: bool, + // {metadata key: {metadata value: offset ids}} + data: HashMap>>, + } + + impl ReferenceState { + fn new() -> Self { + ReferenceState { + in_transaction: false, + data: HashMap::new(), + } + } + + fn key_from_random_number(self: &Self, r1: usize) -> Option { + let keys = self.data.keys().collect::>(); + if keys.len() == 0 { + return None; + } + Some(keys[r1 % keys.len()].clone()) + } + + fn key_and_value_from_random_numbers( + self: &Self, + r1: usize, + r2: usize, + ) -> Option<(String, T)> { + let k = self.key_from_random_number(r1)?; + + let values = self.data.get(&k)?.keys().collect::>(); + if values.len() == 0 { + return None; + } + let v = values[r2 % values.len()]; + Some((k.clone(), v.clone())) + } + + fn key_and_value_and_entry_from_random_numbers( + self: &Self, + r1: usize, + r2: usize, + r3: usize, + ) -> Option<(String, T, u32)> { + let (k, v) = self.key_and_value_from_random_numbers(r1, r2)?; + + let offsets = self.data.get(&k)?.get(&v)?; + if offsets.len() == 0 { + return None; + } + let oid = offsets[r3 % offsets.len()]; + Some((k.clone(), v.clone(), oid)) + } + + fn kv_rbm_eq(self: &Self, rbm: &RoaringBitmap, k: &str, v: &T) -> bool { + match self.data.get(k) { + Some(vv) => match vv.get(v) { + Some(rbm2) => vec_rbm_eq(rbm2, rbm), + None => rbm.is_empty(), + }, + None => rbm.is_empty(), + } + } + } + + pub(crate) struct ReferenceStateMachineImpl { + value_type: PhantomData, + } + + impl ReferenceStateMachine for ReferenceStateMachineImpl { + type State = ReferenceState; + type Transition = Transition; + + fn init_state() -> BoxedStrategy { + return Just(ReferenceState::::new()).boxed(); + } + + fn transitions(state: &ReferenceState) -> BoxedStrategy> { + let r = state.key_and_value_and_entry_from_random_numbers(0, 0, 0); + if r.is_none() { + // Nothing in the reference state yet. + return prop_oneof![ + Just(Transition::BeginTransaction), + Just(Transition::CommitTransaction), + // Set, get, delete random values + (".{0,10}", T::strategy(), 1..1000).prop_map(move |(k, v, oid)| { + Transition::Set(k.to_string(), v, oid as u32) + }), + (".{0,10}", T::strategy(), 1..1000).prop_map(move |(k, v, oid)| { + Transition::Delete(k.to_string(), v, oid as u32) + }), + (".{0,10}", T::strategy()) + .prop_map(move |(k, v)| { Transition::Get(k.to_string(), v) }), + ] + .boxed(); + } + + let state = Rc::new(state.clone()); + return prop_oneof![ + Just(Transition::BeginTransaction), + Just(Transition::CommitTransaction), + // Set, get, delete random values + (".{0,10}", T::strategy(), 1..1000).prop_map(move |(k, v, oid)| { + Transition::Set(k.to_string(), v, oid as u32) + }), + (".{0,10}", T::strategy(), 1..1000).prop_map(move |(k, v, oid)| { + Transition::Delete(k.to_string(), v, oid as u32) + }), + (".{0,10}", T::strategy()).prop_map(move |(k, v)| { + Transition::Get(k.to_string(), v) + }), + + // Sets on values in the reference state + (0..usize::MAX, T::strategy(), 1..1000).prop_map({ + let state = Rc::clone(&state); + move |(r1, v, oid)| { + match state.key_from_random_number(r1) { + Some(k) => Transition::Set(k, v, oid as u32), + None => panic!("Error in the test harness: Setting on key from ref state") + } + } + }), + (0..usize::MAX, 0..usize::MAX, 1..1000).prop_map({ + let state = Rc::clone(&state); + move |(r1, r2, oid)| { + match state.key_and_value_from_random_numbers(r1, r2) { + Some((k, v)) => Transition::Set(k, v, oid as u32), + None => panic!("Error in the test harness: Setting on key and value from ref state") + } + } + }), + (0..usize::MAX, 0..usize::MAX, 0..usize::MAX).prop_map({ + let state = Rc::clone(&state); + move |(r1, r2, r3)| { + match state.key_and_value_and_entry_from_random_numbers(r1, r2, r3) { + Some((k, v, oid)) => Transition::Set(k, v, oid), + None => panic!("Error in the test harness: Setting on key, value, and entry from ref state") + } + } + }), + + // Gets on values in the reference state + (0..usize::MAX, T::strategy()).prop_map({ + let state = Rc::clone(&state); + move |(r1, v)| { + match state.key_from_random_number(r1) { + Some(k) => Transition::Get(k, v), + None => panic!("Error in the test harness: Getting on key from ref state") + } + } + }), + (0..usize::MAX, 0..usize::MAX).prop_map({ + let state = Rc::clone(&state); + move |(r1, r2)| { + match state.key_and_value_from_random_numbers(r1, r2) { + Some((k, v)) => Transition::Get(k, v), + None => panic!("Error in the test harness: Getting on key and value from ref state") + } + } + }), + + // Deletes on values in the reference state + (0..usize::MAX, T::strategy(), 1..1000).prop_map({ + let state = Rc::clone(&state); + move |(r1, v, oid)| { + match state.key_from_random_number(r1) { + Some(k) => Transition::Delete(k, v, oid as u32), + None => panic!("Error in the test harness: Deleting on key from ref state") + } + } + }), + (0..usize::MAX, 0..usize::MAX, 1..1000).prop_map({ + let state = Rc::clone(&state); + move |(r1, r2, oid)| { + match state.key_and_value_from_random_numbers(r1, r2) { + Some((k, v)) => Transition::Delete(k, v, oid as u32), + None => panic!("Error in the test harness: Deleting on key and value from ref state") + } + } + }), + (0..usize::MAX, 0..usize::MAX, 0..usize::MAX).prop_map({ + let state = Rc::clone(&state); + move |(r1, r2, r3)| { + match state.key_and_value_and_entry_from_random_numbers(r1, r2, r3) { + Some((k, v, oid)) => Transition::Delete(k, v, oid as u32), + None => panic!("Error in the test harness: Deleting on key, value, and entry from ref state") + } + } + }), + ].boxed(); + } + + fn apply(mut state: ReferenceState, transition: &Transition) -> Self::State { + match transition { + Transition::BeginTransaction => { + state.in_transaction = true; + } + Transition::CommitTransaction => { + state.in_transaction = false; + } + Transition::Set(k, v, oid) => { + if !state.in_transaction { + return state; + } + let entry = state.data.entry(k.clone()).or_insert(HashMap::new()); + let entry = entry.entry(v.clone()).or_insert(Vec::new()); + if !entry.contains(oid) { + entry.push(*oid); + } + } + Transition::Delete(k, v, oid) => { + if !state.in_transaction { + return state; + } + if !state.data.contains_key(k) { + return state; + } + let entry = state.data.get_mut(k).unwrap(); + if let Some(offsets) = entry.get_mut(v) { + offsets.retain(|x| *x != *oid); + } + + // Remove the offset ID list if it's empty. + let offsets_count = entry.get(v).map(|x| x.len()).unwrap_or(0); + if offsets_count == 0 { + entry.remove(v); + } + + // Remove the entire hashmap for the key if it's empty. + if entry.is_empty() { + state.data.remove(k); + } + } + Transition::Get(_, _) => { + // No-op + } + } + state + } + } + + impl StateMachineTest for BlockfileMetadataIndex { + type SystemUnderTest = Self; + type Reference = ReferenceStateMachineImpl; + + fn init_test(_ref_state: &ReferenceState) -> Self::SystemUnderTest { + // We don't need to set up on _ref_state since we always initialize + // ref_state to empty. + let mut provider = HashMapBlockfileProvider::new(); + let blockfile = provider + .create("test", KeyType::String, ValueType::RoaringBitmap) + .unwrap(); + return BlockfileMetadataIndex::::new(blockfile); + } + + fn apply( + mut state: Self::SystemUnderTest, + ref_state: &ReferenceState, + transition: Transition, + ) -> Self::SystemUnderTest { + match transition { + Transition::BeginTransaction => { + let already_in_transaction = state.in_transaction(); + let res = state.begin_transaction(); + assert!(state.in_transaction()); + if already_in_transaction { + assert!(res.is_err()); + } else { + assert!(res.is_ok()); + } + } + Transition::CommitTransaction => { + let in_transaction = state.in_transaction(); + let res = state.commit_transaction(); + assert_eq!(state.in_transaction(), false); + if !in_transaction { + assert!(res.is_err()); + } else { + assert!(res.is_ok()); + } + } + Transition::Set(k, v, oid) => { + let in_transaction = state.in_transaction(); + let res = state.set(&k, v.clone(), oid); + if !in_transaction { + assert!(res.is_err()); + } else { + assert!(res.is_ok()); + } + } + Transition::Delete(k, v, oid) => { + let in_transaction = state.in_transaction(); + let res = state.delete(&k, v, oid); + if !in_transaction { + assert!(res.is_err()); + } else { + assert!(res.is_ok()); + } + } + Transition::Get(k, v) => { + let in_transaction = state.in_transaction(); + let res = state.get(&k, v.clone()); + if in_transaction { + assert!(res.is_err()); + return state; + } + let rbm = res.unwrap(); + assert!(ref_state.kv_rbm_eq(&rbm, &k, &v)); + } + } + state + } + + fn check_invariants(state: &Self::SystemUnderTest, ref_state: &ReferenceState) { + assert_eq!(state.in_transaction(), ref_state.in_transaction); + if state.in_transaction() { + return; + } + for (metadata_key, metadata_values_hm) in &ref_state.data { + for (metadata_value, ref_offset_ids) in metadata_values_hm { + assert!(vec_rbm_eq( + ref_offset_ids, + &state.get(metadata_key, metadata_value.clone()).unwrap() + )); + } + } + // TODO once we have a way to iterate over all state in the blockfile, + // add a similar check here to make sure that blockfile data is a + // subset of reference state data. + } + } + + prop_state_machine! { + #![proptest_config(Config { + // Enable verbose mode to make the state machine test print the + // transitions for each case. + verbose: 0, + cases: 100, + .. Config::default() + })] + #[test] + fn proptest_string_metadata_index(sequential 1..100 => BlockfileMetadataIndex); + + #[test] + fn proptest_boolean_metadata_index(sequential 1..100 => BlockfileMetadataIndex); + + #[test] + fn proptest_numeric_metadata_index(sequential 1..100 => BlockfileMetadataIndex); + } +} diff --git a/rust/worker/src/index/mod.rs b/rust/worker/src/index/mod.rs index ddaf8d737a4..6708cfb4547 100644 --- a/rust/worker/src/index/mod.rs +++ b/rust/worker/src/index/mod.rs @@ -1,7 +1,11 @@ +mod fulltext; mod hnsw; +mod metadata; mod types; mod utils; // Re-export types + pub(crate) use hnsw::*; + pub(crate) use types::*; diff --git a/rust/worker/src/index/types.rs b/rust/worker/src/index/types.rs index 7af440c947c..092e3fa0c45 100644 --- a/rust/worker/src/index/types.rs +++ b/rust/worker/src/index/types.rs @@ -1,3 +1,4 @@ +use crate::distance::DistanceFunction; use crate::errors::{ChromaError, ErrorCodes}; use crate::types::{MetadataValue, Segment}; use thiserror::Error; @@ -78,58 +79,3 @@ pub(crate) trait PersistentIndex: Index { where Self: Sized; } - -/// The distance function enum. -/// # Description -/// This enum defines the distance functions supported by indices in Chroma. -/// # Variants -/// - `Euclidean` - The Euclidean or l2 norm. -/// - `Cosine` - The cosine distance. Specifically, 1 - cosine. -/// - `InnerProduct` - The inner product. Specifically, 1 - inner product. -/// # Notes -/// See https://docs.trychroma.com/usage-guide#changing-the-distance-function -#[derive(Clone, Debug)] -pub(crate) enum DistanceFunction { - Euclidean, - Cosine, - InnerProduct, -} - -#[derive(Error, Debug)] -pub(crate) enum DistanceFunctionError { - #[error("Invalid distance function `{0}`")] - InvalidDistanceFunction(String), -} - -impl ChromaError for DistanceFunctionError { - fn code(&self) -> ErrorCodes { - match self { - DistanceFunctionError::InvalidDistanceFunction(_) => ErrorCodes::InvalidArgument, - } - } -} - -impl TryFrom<&str> for DistanceFunction { - type Error = DistanceFunctionError; - - fn try_from(value: &str) -> Result { - match value { - "l2" => Ok(DistanceFunction::Euclidean), - "cosine" => Ok(DistanceFunction::Cosine), - "ip" => Ok(DistanceFunction::InnerProduct), - _ => Err(DistanceFunctionError::InvalidDistanceFunction( - value.to_string(), - )), - } - } -} - -impl Into for DistanceFunction { - fn into(self) -> String { - match self { - DistanceFunction::Euclidean => "l2".to_string(), - DistanceFunction::Cosine => "cosine".to_string(), - DistanceFunction::InnerProduct => "ip".to_string(), - } - } -} diff --git a/rust/worker/src/ingest/config.rs b/rust/worker/src/ingest/config.rs deleted file mode 100644 index b7647cfe30e..00000000000 --- a/rust/worker/src/ingest/config.rs +++ /dev/null @@ -1,6 +0,0 @@ -use serde::Deserialize; - -#[derive(Deserialize)] -pub(crate) struct IngestConfig { - pub(crate) queue_size: usize, -} diff --git a/rust/worker/src/ingest/ingest.rs b/rust/worker/src/ingest/ingest.rs deleted file mode 100644 index bacf627cb76..00000000000 --- a/rust/worker/src/ingest/ingest.rs +++ /dev/null @@ -1,417 +0,0 @@ -use async_trait::async_trait; -use bytes::Bytes; -use futures::{StreamExt, TryStreamExt}; -use prost::Message; -use std::{ - collections::{HashMap, HashSet}, - fmt::Debug, - sync::{Arc, RwLock}, -}; - -use crate::{ - assignment::{ - self, - assignment_policy::{self, AssignmentPolicy}, - }, - chroma_proto, - config::{Configurable, WorkerConfig}, - errors::{ChromaError, ErrorCodes}, - memberlist::{CustomResourceMemberlistProvider, Memberlist}, - sysdb::sysdb::{GrpcSysDb, SysDb}, - system::{Component, ComponentContext, ComponentHandle, Handler, Receiver, StreamHandler}, - types::{EmbeddingRecord, EmbeddingRecordConversionError, SeqId}, -}; - -use pulsar::{Consumer, DeserializeMessage, Payload, Pulsar, SubType, TokioExecutor}; -use thiserror::Error; - -use super::message_id::PulsarMessageIdWrapper; - -/// An ingest component is responsible for ingesting data into the system from the log -/// stream. -/// # Notes -/// The only current implementation of the ingest is the Pulsar ingest. -pub(crate) struct Ingest { - assignment_policy: RwLock>, - assigned_topics: RwLock>, - topic_to_handle: RwLock>>, - queue_size: usize, - my_ip: String, - pulsar_tenant: String, - pulsar_namespace: String, - pulsar: Pulsar, - sysdb: Box, - scheduler: Option)>>>, -} - -impl Component for Ingest { - fn queue_size(&self) -> usize { - return self.queue_size; - } -} - -impl Debug for Ingest { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Ingest") - .field("queue_size", &self.queue_size) - .finish() - } -} - -#[derive(Error, Debug)] -pub(crate) enum IngestConfigurationError { - #[error(transparent)] - PulsarError(#[from] pulsar::Error), -} - -impl ChromaError for IngestConfigurationError { - fn code(&self) -> ErrorCodes { - match self { - IngestConfigurationError::PulsarError(_e) => ErrorCodes::Internal, - } - } -} - -// TODO: Nest the ingest assignment policy inside the ingest component config so its -// specific to the ingest component and can be used here -#[async_trait] -impl Configurable for Ingest { - async fn try_from_config(worker_config: &WorkerConfig) -> Result> { - let assignment_policy = assignment_policy::RendezvousHashingAssignmentPolicy::new( - worker_config.pulsar_tenant.clone(), - worker_config.pulsar_namespace.clone(), - ); - - println!("Pulsar connection url: {}", worker_config.pulsar_url); - let pulsar = match Pulsar::builder(worker_config.pulsar_url.clone(), TokioExecutor) - .build() - .await - { - Ok(pulsar) => pulsar, - Err(e) => { - return Err(Box::new(IngestConfigurationError::PulsarError(e))); - } - }; - - // TODO: Sysdb should have a dynamic resolution in sysdb - let sysdb = GrpcSysDb::try_from_config(worker_config).await; - let sysdb = match sysdb { - Ok(sysdb) => sysdb, - Err(err) => { - return Err(err); - } - }; - - let ingest = Ingest { - assignment_policy: RwLock::new(Box::new(assignment_policy)), - assigned_topics: RwLock::new(vec![]), - topic_to_handle: RwLock::new(HashMap::new()), - queue_size: worker_config.ingest.queue_size, - my_ip: worker_config.my_ip.clone(), - pulsar: pulsar, - pulsar_tenant: worker_config.pulsar_tenant.clone(), - pulsar_namespace: worker_config.pulsar_namespace.clone(), - sysdb: Box::new(sysdb), - scheduler: None, - }; - Ok(ingest) - } -} - -impl Ingest { - fn get_topics(&self) -> Vec { - // This mirrors the current python and go code, which assumes a fixed set of topics - let mut topics = Vec::with_capacity(16); - for i in 0..16 { - let topic = format!( - "persistent://{}/{}/chroma_log_{}", - self.pulsar_tenant, self.pulsar_namespace, i - ); - topics.push(topic); - } - return topics; - } - - pub(crate) fn subscribe( - &mut self, - scheduler: Box)>>, - ) { - self.scheduler = Some(scheduler); - } -} - -#[async_trait] -impl Handler for Ingest { - async fn handle(&mut self, msg: Memberlist, ctx: &ComponentContext) { - let mut new_assignments = HashSet::new(); - let candidate_topics: Vec = self.get_topics(); - println!( - "Performing assignment for topics: {:?}. My ip: {}", - candidate_topics, self.my_ip - ); - // Scope for assigner write lock to be released so we don't hold it over await - { - let mut assigner = match self.assignment_policy.write() { - Ok(assigner) => assigner, - Err(err) => { - println!("Failed to read assignment policy: {:?}", err); - return; - } - }; - - // Use the assignment policy to assign topics to this worker - assigner.set_members(msg); - for topic in candidate_topics.iter() { - let assignment = assigner.assign(topic); - let assignment = match assignment { - Ok(assignment) => assignment, - Err(err) => { - // TODO: Log error - continue; - } - }; - if assignment == self.my_ip { - new_assignments.insert(topic); - } - } - } - - // Compute the topics we need to add/remove - let mut to_remove = Vec::new(); - let mut to_add = Vec::new(); - - // Scope for assigned topics read lock to be released so we don't hold it over await - { - let assigned_topics_handle = self.assigned_topics.read(); - match assigned_topics_handle { - Ok(assigned_topics) => { - // Compute the diff between the current assignments and the new assignments - for topic in assigned_topics.iter() { - if !new_assignments.contains(topic) { - to_remove.push(topic.to_string()); - } - } - for topic in new_assignments.iter() { - if !assigned_topics.contains(*topic) { - to_add.push(topic.to_string()); - } - } - } - Err(err) => { - // TODO: Log error and handle lock poisoning - } - } - } - - // Unsubscribe from topics we no longer need to listen to - for topic in to_remove.iter() { - match self.topic_to_handle.write() { - Ok(mut topic_to_handle) => { - let handle = topic_to_handle.remove(topic); - match handle { - Some(mut handle) => { - handle.stop(); - } - None => { - // TODO: This should log an error - println!("No handle found for topic: {}", topic); - } - } - } - Err(err) => { - // TODO: Log an error and handle lock poisoning - } - } - } - - // Subscribe to new topics - for topic in to_add.iter() { - // Do the subscription and register the stream to this ingest component - let consumer: Consumer = self - .pulsar - .consumer() - .with_topic(topic.to_string()) - .with_subscription_type(SubType::Exclusive) - .build() - .await - .unwrap(); - println!("Created consumer for topic: {}", topic); - - let scheduler = match &self.scheduler { - Some(scheduler) => scheduler.clone(), - None => { - // TODO: log error - return; - } - }; - - let ingest_topic_component = - PulsarIngestTopic::new(consumer, self.sysdb.clone(), scheduler); - - let handle = ctx.system.clone().start_component(ingest_topic_component); - - // Bookkeep the handle so we can shut the stream down later - match self.topic_to_handle.write() { - Ok(mut topic_to_handle) => { - topic_to_handle.insert(topic.to_string(), handle); - } - Err(err) => { - // TODO: log error and handle lock poisoning - println!("Failed to write topic to handle: {:?}", err); - } - } - } - } -} - -impl DeserializeMessage for chroma_proto::SubmitEmbeddingRecord { - type Output = Self; - - fn deserialize_message(payload: &Payload) -> chroma_proto::SubmitEmbeddingRecord { - // Its a bit strange to unwrap here, but the pulsar api doesn't give us a way to - // return an error, so we have to panic if we can't decode the message - // also we are forced to clone since the api doesn't give us a way to borrow the bytes - // TODO: can we not clone? - // TODO: I think just typing this to Result<> would allow errors to propagate - let record = - chroma_proto::SubmitEmbeddingRecord::decode(Bytes::from(payload.data.clone())).unwrap(); - return record; - } -} - -struct PulsarIngestTopic { - consumer: RwLock>>, - sysdb: Box, - scheduler: Box)>>, -} - -impl Debug for PulsarIngestTopic { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PulsarIngestTopic").finish() - } -} - -impl PulsarIngestTopic { - fn new( - consumer: Consumer, - sysdb: Box, - scheduler: Box)>>, - ) -> Self { - PulsarIngestTopic { - consumer: RwLock::new(Some(consumer)), - sysdb: sysdb, - scheduler: scheduler, - } - } -} - -impl Component for PulsarIngestTopic { - fn queue_size(&self) -> usize { - 1000 - } - - fn on_start(&mut self, ctx: &ComponentContext) -> () { - println!("Starting PulsarIngestTopic for topic"); - let stream = match self.consumer.write() { - Ok(mut consumer_handle) => consumer_handle.take(), - Err(err) => { - println!("Failed to take consumer handle: {:?}", err); - None - } - }; - let stream = match stream { - Some(stream) => stream, - None => { - return; - } - }; - let stream = stream.then(|result| async { - match result { - Ok(msg) => { - println!( - "PulsarIngestTopic received message with id: {:?}", - msg.message_id - ); - // Convert the Pulsar Message to an EmbeddingRecord - let proto_embedding_record = msg.deserialize(); - let id = msg.message_id; - let seq_id: SeqId = PulsarMessageIdWrapper(id).into(); - let embedding_record: Result = - (proto_embedding_record, seq_id).try_into(); - match embedding_record { - Ok(embedding_record) => { - return Some(Box::new(embedding_record)); - } - Err(err) => { - // TODO: Handle and log - println!("PulsarIngestTopic received error while performing conversion: {:?}", err); - } - } - None - } - Err(err) => { - // TODO: Log an error - println!("PulsarIngestTopic received error: {:?}", err); - // Put this on a dead letter queue, this concept does not exist in our - // system yet - None - } - } - }); - self.register_stream(stream, ctx); - } -} - -#[async_trait] -impl Handler>> for PulsarIngestTopic { - async fn handle( - &mut self, - message: Option>, - _ctx: &ComponentContext, - ) -> () { - // Use the sysdb to tenant id for the embedding record - let embedding_record = match message { - Some(embedding_record) => embedding_record, - None => { - return; - } - }; - - // TODO: Cache this - let coll = self - .sysdb - .get_collections(Some(embedding_record.collection_id), None, None, None, None) - .await; - - let coll = match coll { - Ok(coll) => coll, - Err(err) => { - println!( - "PulsarIngestTopic received error while fetching collection: {:?}", - err - ); - return; - } - }; - - let coll = match coll.first() { - Some(coll) => coll, - None => { - println!("PulsarIngestTopic received empty collection"); - return; - } - }; - - let tenant_id = &coll.tenant; - - let _ = self - .scheduler - .send((tenant_id.clone(), embedding_record)) - .await; - - // TODO: Handle res - } -} - -#[async_trait] -impl StreamHandler>> for PulsarIngestTopic {} diff --git a/rust/worker/src/ingest/message_id.rs b/rust/worker/src/ingest/message_id.rs deleted file mode 100644 index 3ac3d05a1ea..00000000000 --- a/rust/worker/src/ingest/message_id.rs +++ /dev/null @@ -1,48 +0,0 @@ -use std::ops::Deref; - -// mirrors chromadb/utils/messageid.py -use num_bigint::BigInt; -use pulsar::{consumer::data::MessageData, proto::MessageIdData}; - -use crate::types::SeqId; - -pub(crate) struct PulsarMessageIdWrapper(pub(crate) MessageData); - -impl Deref for PulsarMessageIdWrapper { - type Target = MessageIdData; - - fn deref(&self) -> &Self::Target { - &self.0.id - } -} - -pub(crate) fn pulsar_to_int(message_id: PulsarMessageIdWrapper) -> SeqId { - let ledger_id = message_id.ledger_id; - let entry_id = message_id.entry_id; - let batch_index = message_id.batch_index.unwrap_or(0); - let partition = message_id.partition.unwrap_or(0); - - let mut ledger_id = BigInt::from(ledger_id); - let mut entry_id = BigInt::from(entry_id); - let mut batch_index = BigInt::from(batch_index); - let mut partition = BigInt::from(partition); - - // Convert to offset binary encoding to preserve ordering semantics when encoded - // see https://en.wikipedia.org/wiki/Offset_binary - ledger_id = ledger_id + BigInt::from(2).pow(63); - entry_id = entry_id + BigInt::from(2).pow(63); - batch_index = batch_index + BigInt::from(2).pow(31); - partition = partition + BigInt::from(2).pow(31); - - let res = ledger_id << 128 | entry_id << 96 | batch_index << 64 | partition; - res -} - -// We can't use From because we don't own the type -// So the pattern is to wrap it in a newtype and implement TryFrom for that -// And implement Dereference for the newtype to the underlying type -impl From for SeqId { - fn from(message_id: PulsarMessageIdWrapper) -> Self { - return pulsar_to_int(message_id); - } -} diff --git a/rust/worker/src/ingest/mod.rs b/rust/worker/src/ingest/mod.rs deleted file mode 100644 index ae7aaf8d7b5..00000000000 --- a/rust/worker/src/ingest/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -pub(crate) mod config; -mod ingest; -mod message_id; -mod scheduler; - -// Re-export the ingest provider for use in the worker -pub(crate) use ingest::*; -pub(crate) use scheduler::*; diff --git a/rust/worker/src/ingest/scheduler.rs b/rust/worker/src/ingest/scheduler.rs deleted file mode 100644 index 770e9bb0bbf..00000000000 --- a/rust/worker/src/ingest/scheduler.rs +++ /dev/null @@ -1,212 +0,0 @@ -// A scheduler recieves embedding records for a given batch of documents -// and schedules them to be ingested to the segment manager - -use crate::{ - system::{Component, ComponentContext, Handler, Receiver}, - types::EmbeddingRecord, -}; -use async_trait::async_trait; -use rand::prelude::SliceRandom; -use rand::Rng; -use std::{ - collections::{btree_map::Range, HashMap}, - fmt::{Debug, Formatter, Result}, - sync::Arc, -}; - -pub(crate) struct RoundRobinScheduler { - // The segment manager to schedule to, a segment manager is a component - // segment_manager: SegmentManager - curr_wake_up: Option>, - tenant_to_queue: HashMap>>, - new_tenant_channel: Option>, - subscribers: Option>>>>, -} - -impl Debug for RoundRobinScheduler { - fn fmt(&self, f: &mut Formatter<'_>) -> Result { - f.debug_struct("Scheduler").finish() - } -} - -impl RoundRobinScheduler { - pub(crate) fn new() -> Self { - RoundRobinScheduler { - curr_wake_up: None, - tenant_to_queue: HashMap::new(), - new_tenant_channel: None, - subscribers: Some(Vec::new()), - } - } - - pub(crate) fn subscribe(&mut self, subscriber: Box>>) { - match self.subscribers { - Some(ref mut subscribers) => { - subscribers.push(subscriber); - } - None => {} - } - } -} - -impl Component for RoundRobinScheduler { - fn queue_size(&self) -> usize { - 1000 - } - - fn on_start(&mut self, ctx: &ComponentContext) { - let sleep_sender = ctx.sender.clone(); - let (new_tenant_tx, mut new_tenant_rx) = tokio::sync::mpsc::channel(1000); - self.new_tenant_channel = Some(new_tenant_tx); - let cancellation_token = ctx.cancellation_token.clone(); - let subscribers = self.subscribers.take(); - let mut subscribers = match subscribers { - Some(subscribers) => subscribers, - None => { - // TODO: log + error - return; - } - }; - tokio::spawn(async move { - let mut tenant_queues: HashMap< - String, - tokio::sync::mpsc::Receiver>, - > = HashMap::new(); - loop { - // TODO: handle cancellation - let mut did_work = false; - for tenant_queue in tenant_queues.values_mut() { - match tenant_queue.try_recv() { - Ok(message) => { - // Randomly pick a subscriber to send the message to - // This serves as a crude load balancing between available threads - // Future improvements here could be - // - Use a work stealing scheduler - // - Use rayon - // - We need to enforce partial order over writes to a given key - // so we need a mechanism to ensure that all writes to a given key - // occur in order - let mut subscriber = None; - { - let mut rng = rand::thread_rng(); - subscriber = subscribers.choose_mut(&mut rng); - } - match subscriber { - Some(subscriber) => { - let res = subscriber.send(message).await; - } - None => {} - } - did_work = true; - } - Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { - continue; - } - Err(_) => { - // TODO: Handle a erroneous channel - // log an error - continue; - } - }; - } - - match new_tenant_rx.try_recv() { - Ok(new_tenant_message) => { - tenant_queues.insert(new_tenant_message.tenant, new_tenant_message.channel); - } - Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { - // no - op - } - Err(_) => { - // TODO: handle erroneous channel - // log an error - continue; - } - }; - - if !did_work { - // Send a sleep message to the sender - let (wake_tx, wake_rx) = tokio::sync::oneshot::channel(); - let sleep_res = sleep_sender.send(SleepMessage { sender: wake_tx }).await; - let wake_res = wake_rx.await; - } - } - }); - } -} - -#[async_trait] -impl Handler<(String, Box)> for RoundRobinScheduler { - async fn handle( - &mut self, - message: (String, Box), - _ctx: &ComponentContext, - ) { - let (tenant, embedding_record) = message; - // Check if the tenant is already in the tenant set, if not we need to inform the scheduler loop - // of a new tenant - if self.tenant_to_queue.get(&tenant).is_none() { - // Create a new channel for the tenant - let (sender, reciever) = tokio::sync::mpsc::channel(1000); - // Add the tenant to the tenant set - self.tenant_to_queue.insert(tenant.clone(), sender); - // Send the new tenant message to the scheduler loop - let new_tenant_channel = match self.new_tenant_channel { - Some(ref mut channel) => channel, - None => { - // TODO: this is an error - // It should always be populated by on_start - return; - } - }; - let res = new_tenant_channel - .send(NewTenantMessage { - tenant: tenant.clone(), - channel: reciever, - }) - .await; - // TODO: handle this res - } - - // Send the embedding record to the tenant's channel - let res = self - .tenant_to_queue - .get(&tenant) - .unwrap() - .send(embedding_record) - .await; - // TODO: handle this res - - // Check if the scheduler is sleeping, if so wake it up - // TODO: we need to init with a wakeup otherwise we are off by one - if self.curr_wake_up.is_some() { - // Send a wake up message to the scheduler loop - let res = self.curr_wake_up.take().unwrap().send(WakeMessage {}); - // TOOD: handle this res - } - } -} - -#[async_trait] -impl Handler for RoundRobinScheduler { - async fn handle(&mut self, message: SleepMessage, _ctx: &ComponentContext) { - // Set the current wake up channel - self.curr_wake_up = Some(message.sender); - } -} - -/// Used by round robin scheduler to wake its scheduler loop -#[derive(Debug)] -struct WakeMessage {} - -/// The round robin scheduler will sleep when there is no work to be done and send a sleep message -/// this allows the manager to wake it up when there is work to be scheduled -#[derive(Debug)] -struct SleepMessage { - sender: tokio::sync::oneshot::Sender, -} - -struct NewTenantMessage { - tenant: String, - channel: tokio::sync::mpsc::Receiver>, -} diff --git a/rust/worker/src/lib.rs b/rust/worker/src/lib.rs index ae7ea7dc7d5..4f89e95d78f 100644 --- a/rust/worker/src/lib.rs +++ b/rust/worker/src/lib.rs @@ -1,8 +1,12 @@ mod assignment; +mod blockstore; +mod compactor; mod config; +mod distance; mod errors; +mod execution; mod index; -mod ingest; +mod log; mod memberlist; mod segment; mod server; @@ -14,90 +18,85 @@ mod types; use config::Configurable; use memberlist::MemberlistProvider; -use crate::sysdb::sysdb::SysDb; - mod chroma_proto { tonic::include_proto!("chroma"); } -pub async fn worker_entrypoint() { +pub async fn query_service_entrypoint() { let config = config::RootConfig::load(); - // Create all the core components and start them - // TODO: This should be handled by an Application struct and we can push the config into it - // for now we expose the config to pub and inject it into the components - - // The two root components are ingest, and the gRPC server - let mut system: system::System = system::System::new(); - - let mut ingest = match ingest::Ingest::try_from_config(&config.worker).await { - Ok(ingest) => ingest, - Err(err) => { - println!("Failed to create ingest component: {:?}", err); - return; - } - }; - - let mut memberlist = - match memberlist::CustomResourceMemberlistProvider::try_from_config(&config.worker).await { - Ok(memberlist) => memberlist, + let config = config.query_service; + let system: system::System = system::System::new(); + let dispatcher = + match execution::dispatcher::Dispatcher::try_from_config(&config.dispatcher).await { + Ok(dispatcher) => dispatcher, Err(err) => { - println!("Failed to create memberlist component: {:?}", err); + println!("Failed to create dispatcher component: {:?}", err); return; } }; - - let mut scheduler = ingest::RoundRobinScheduler::new(); - - let segment_manager = match segment::SegmentManager::try_from_config(&config.worker).await { - Ok(segment_manager) => segment_manager, + let mut dispatcher_handle = system.start_component(dispatcher); + let mut worker_server = match server::WorkerServer::try_from_config(&config).await { + Ok(worker_server) => worker_server, Err(err) => { - println!("Failed to create segment manager component: {:?}", err); + println!("Failed to create worker server component: {:?}", err); return; } }; + worker_server.set_system(system); + worker_server.set_dispatcher(dispatcher_handle.receiver()); - let mut segment_ingestor_receivers = - Vec::with_capacity(config.worker.num_indexing_threads as usize); - for _ in 0..config.worker.num_indexing_threads { - let segment_ingestor = segment::SegmentIngestor::new(segment_manager.clone()); - let segment_ingestor_handle = system.start_component(segment_ingestor); - let recv = segment_ingestor_handle.receiver(); - segment_ingestor_receivers.push(recv); - } + let server_join_handle = tokio::spawn(async move { + let _ = crate::server::WorkerServer::run(worker_server).await; + }); - let mut worker_server = match server::WorkerServer::try_from_config(&config.worker).await { - Ok(worker_server) => worker_server, + let _ = tokio::join!(server_join_handle, dispatcher_handle.join()); +} + +pub async fn compaction_service_entrypoint() { + let config = config::RootConfig::load(); + let config = config.compaction_service; + let system: system::System = system::System::new(); + + let mut memberlist = match memberlist::CustomResourceMemberlistProvider::try_from_config( + &config.memberlist_provider, + ) + .await + { + Ok(memberlist) => memberlist, Err(err) => { - println!("Failed to create worker server component: {:?}", err); + println!("Failed to create memberlist component: {:?}", err); return; } }; - worker_server.set_segment_manager(segment_manager.clone()); - - // Boot the system - // memberlist -> ingest -> scheduler -> NUM_THREADS x segment_ingestor -> segment_manager - // server <- segment_manager - for recv in segment_ingestor_receivers { - scheduler.subscribe(recv); - } + let dispatcher = + match execution::dispatcher::Dispatcher::try_from_config(&config.dispatcher).await { + Ok(dispatcher) => dispatcher, + Err(err) => { + println!("Failed to create dispatcher component: {:?}", err); + return; + } + }; + let mut dispatcher_handle = system.start_component(dispatcher); + let mut compaction_manager = + match crate::compactor::CompactionManager::try_from_config(&config).await { + Ok(compaction_manager) => compaction_manager, + Err(err) => { + println!("Failed to create compaction manager component: {:?}", err); + return; + } + }; + compaction_manager.set_dispatcher(dispatcher_handle.receiver()); + compaction_manager.set_system(system.clone()); - let mut scheduler_handler = system.start_component(scheduler); - ingest.subscribe(scheduler_handler.receiver()); + let mut compaction_manager_handle = system.start_component(compaction_manager); + memberlist.subscribe(compaction_manager_handle.receiver()); - let mut ingest_handle = system.start_component(ingest); - let recv = ingest_handle.receiver(); - memberlist.subscribe(recv); let mut memberlist_handle = system.start_component(memberlist); - let server_join_handle = tokio::spawn(async move { - crate::server::WorkerServer::run(worker_server).await; - }); - - // Join on all handles - let _ = tokio::join!( - ingest_handle.join(), + tokio::join!( memberlist_handle.join(), - scheduler_handler.join(), + compaction_manager_handle.join(), + dispatcher_handle.join() ); } diff --git a/rust/worker/src/log/config.rs b/rust/worker/src/log/config.rs new file mode 100644 index 00000000000..ff9196bfdfa --- /dev/null +++ b/rust/worker/src/log/config.rs @@ -0,0 +1,12 @@ +use serde::Deserialize; + +#[derive(Deserialize)] +pub(crate) struct GrpcLogConfig { + pub(crate) host: String, + pub(crate) port: u16, +} + +#[derive(Deserialize)] +pub(crate) enum LogConfig { + Grpc(GrpcLogConfig), +} diff --git a/rust/worker/src/log/log.rs b/rust/worker/src/log/log.rs new file mode 100644 index 00000000000..0b69a81f532 --- /dev/null +++ b/rust/worker/src/log/log.rs @@ -0,0 +1,313 @@ +use crate::chroma_proto; +use crate::chroma_proto::log_service_client::LogServiceClient; +use crate::config::Configurable; +use crate::errors::ChromaError; +use crate::errors::ErrorCodes; +use crate::log::config::LogConfig; +use crate::types::LogRecord; +use crate::types::RecordConversionError; +use async_trait::async_trait; +use std::collections::HashMap; +use std::fmt::Debug; +use thiserror::Error; +use uuid::Uuid; + +/// CollectionInfo is a struct that contains information about a collection for the +/// compacting process. +/// Fields: +/// - collection_id: the id of the collection that needs to be compacted +/// - first_log_offset: the offset of the first log entry in the collection that needs to be compacted +/// - first_log_ts: the timestamp of the first log entry in the collection that needs to be compacted +pub(crate) struct CollectionInfo { + pub(crate) collection_id: String, + pub(crate) first_log_offset: i64, + pub(crate) first_log_ts: i64, +} + +#[derive(Clone, Debug)] +pub(crate) struct CollectionRecord { + pub(crate) id: String, + pub(crate) tenant_id: String, + pub(crate) last_compaction_time: i64, + pub(crate) first_record_time: i64, + pub(crate) offset: i64, +} + +#[async_trait] +pub(crate) trait Log: Send + Sync + LogClone + Debug { + async fn read( + &mut self, + collection_id: Uuid, + offset: i64, + batch_size: i32, + end_timestamp: Option, + ) -> Result, PullLogsError>; + + async fn get_collections_with_new_data( + &mut self, + ) -> Result, GetCollectionsWithNewDataError>; +} + +pub(crate) trait LogClone { + fn clone_box(&self) -> Box; +} + +impl LogClone for T +where + T: 'static + Log + Clone, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + fn clone(&self) -> Box { + self.clone_box() + } +} + +#[derive(Clone, Debug)] +pub(crate) struct GrpcLog { + client: LogServiceClient, +} + +impl GrpcLog { + pub(crate) fn new(client: LogServiceClient) -> Self { + Self { client } + } +} + +#[derive(Error, Debug)] +pub(crate) enum GrpcLogError { + #[error("Failed to connect to log service")] + FailedToConnect(#[from] tonic::transport::Error), +} + +impl ChromaError for GrpcLogError { + fn code(&self) -> ErrorCodes { + match self { + GrpcLogError::FailedToConnect(_) => ErrorCodes::Internal, + } + } +} + +#[async_trait] +impl Configurable for GrpcLog { + async fn try_from_config(config: &LogConfig) -> Result> { + match &config { + LogConfig::Grpc(my_config) => { + let host = &my_config.host; + let port = &my_config.port; + // TODO: switch to logging when logging is implemented + println!("Connecting to log service at {}:{}", host, port); + let connection_string = format!("http://{}:{}", host, port); + let client = LogServiceClient::connect(connection_string).await; + match client { + Ok(client) => { + return Ok(GrpcLog::new(client)); + } + Err(e) => { + return Err(Box::new(GrpcLogError::FailedToConnect(e))); + } + } + } + } + } +} + +#[async_trait] +impl Log for GrpcLog { + async fn read( + &mut self, + collection_id: Uuid, + offset: i64, + batch_size: i32, + end_timestamp: Option, + ) -> Result, PullLogsError> { + let end_timestamp = match end_timestamp { + Some(end_timestamp) => end_timestamp, + None => -1, + }; + let request = self.client.pull_logs(chroma_proto::PullLogsRequest { + collection_id: collection_id.to_string(), + start_from_offset: offset, + batch_size, + end_timestamp, + }); + let response = request.await; + match response { + Ok(response) => { + let logs = response.into_inner().records; + let mut result = Vec::new(); + for log_record_proto in logs { + let log_record = log_record_proto.try_into(); + match log_record { + Ok(log_record) => { + result.push(log_record); + } + Err(err) => { + return Err(PullLogsError::ConversionError(err)); + } + } + } + Ok(result) + } + Err(e) => { + // TODO: switch to logging when logging is implemented + println!("Failed to pull logs: {}", e); + Err(PullLogsError::FailedToPullLogs(e)) + } + } + } + + async fn get_collections_with_new_data( + &mut self, + ) -> Result, GetCollectionsWithNewDataError> { + let request = self.client.get_all_collection_info_to_compact( + chroma_proto::GetAllCollectionInfoToCompactRequest {}, + ); + let response = request.await; + + match response { + Ok(response) => { + let collections = response.into_inner().all_collection_info; + let mut result = Vec::new(); + for collection in collections { + result.push(CollectionInfo { + collection_id: collection.collection_id, + first_log_offset: collection.first_log_offset, + first_log_ts: collection.first_log_ts, + }); + } + Ok(result) + } + Err(e) => { + // TODO: switch to logging when logging is implemented + println!("Failed to get collections: {}", e); + Err(GetCollectionsWithNewDataError::FailedGetCollectionsWithNewData(e)) + } + } + } +} + +#[derive(Error, Debug)] +pub(crate) enum PullLogsError { + #[error("Failed to fetch")] + FailedToPullLogs(#[from] tonic::Status), + #[error("Failed to convert proto embedding record into EmbeddingRecord")] + ConversionError(#[from] RecordConversionError), +} + +impl ChromaError for PullLogsError { + fn code(&self) -> ErrorCodes { + match self { + PullLogsError::FailedToPullLogs(_) => ErrorCodes::Internal, + PullLogsError::ConversionError(_) => ErrorCodes::Internal, + } + } +} + +#[derive(Error, Debug)] +pub(crate) enum GetCollectionsWithNewDataError { + #[error("Failed to fetch")] + FailedGetCollectionsWithNewData(#[from] tonic::Status), +} + +impl ChromaError for GetCollectionsWithNewDataError { + fn code(&self) -> ErrorCodes { + match self { + GetCollectionsWithNewDataError::FailedGetCollectionsWithNewData(_) => { + ErrorCodes::Internal + } + } + } +} + +// This is used for testing only, it represents a log record that is stored in memory +// internal to a mock log implementation +#[derive(Clone)] +pub(crate) struct InternalLogRecord { + pub(crate) collection_id: String, + pub(crate) log_offset: i64, + pub(crate) log_ts: i64, + pub(crate) record: LogRecord, +} + +impl Debug for InternalLogRecord { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LogRecord") + .field("collection_id", &self.collection_id) + .field("log_offset", &self.log_offset) + .field("log_ts", &self.log_ts) + .field("record", &self.record) + .finish() + } +} + +// This is used for testing only +#[derive(Clone, Debug)] +pub(crate) struct InMemoryLog { + logs: HashMap>>, +} + +impl InMemoryLog { + pub fn new() -> InMemoryLog { + InMemoryLog { + logs: HashMap::new(), + } + } + + pub fn add_log(&mut self, collection_id: String, log: Box) { + let logs = self.logs.entry(collection_id).or_insert(Vec::new()); + logs.push(log); + } +} + +#[async_trait] +impl Log for InMemoryLog { + async fn read( + &mut self, + collection_id: Uuid, + offset: i64, + batch_size: i32, + end_timestamp: Option, + ) -> Result, PullLogsError> { + let end_timestamp = match end_timestamp { + Some(end_timestamp) => end_timestamp, + None => i64::MAX, + }; + + let logs = match self.logs.get(&collection_id.to_string()) { + Some(logs) => logs, + None => return Ok(Vec::new()), + }; + let mut result = Vec::new(); + for i in offset..(offset + batch_size as i64) { + if i < logs.len() as i64 && logs[i as usize].log_ts <= end_timestamp { + result.push(logs[i as usize].record.clone()); + } + } + Ok(result) + } + + async fn get_collections_with_new_data( + &mut self, + ) -> Result, GetCollectionsWithNewDataError> { + let mut collections = Vec::new(); + for (collection_id, log_record) in self.logs.iter() { + if log_record.is_empty() { + continue; + } + // sort the logs by log_offset + let mut logs = log_record.clone(); + logs.sort_by(|a, b| a.log_offset.cmp(&b.log_offset)); + collections.push(CollectionInfo { + collection_id: collection_id.clone(), + first_log_offset: logs[0].log_offset, + first_log_ts: logs[0].log_ts, + }); + } + Ok(collections) + } +} diff --git a/rust/worker/src/log/mod.rs b/rust/worker/src/log/mod.rs new file mode 100644 index 00000000000..f6b76068c83 --- /dev/null +++ b/rust/worker/src/log/mod.rs @@ -0,0 +1,16 @@ +pub(crate) mod config; +pub(crate) mod log; + +use crate::{config::Configurable, errors::ChromaError}; + +use self::config::LogConfig; + +pub(crate) async fn from_config( + config: &LogConfig, +) -> Result, Box> { + match &config { + crate::log::config::LogConfig::Grpc(_) => { + Ok(Box::new(log::GrpcLog::try_from_config(config).await?)) + } + } +} diff --git a/rust/worker/src/memberlist/config.rs b/rust/worker/src/memberlist/config.rs index d6aaf2c8682..f94088941fa 100644 --- a/rust/worker/src/memberlist/config.rs +++ b/rust/worker/src/memberlist/config.rs @@ -18,10 +18,12 @@ pub(crate) enum MemberlistProviderConfig { /// The configuration for the custom resource memberlist provider. /// # Fields +/// - kube_namespace: The namespace to use for the custom resource. /// - memberlist_name: The name of the custom resource to use for the memberlist. /// - queue_size: The size of the queue to use for the channel. #[derive(Deserialize)] pub(crate) struct CustomResourceMemberlistProviderConfig { + pub(crate) kube_namespace: String, pub(crate) memberlist_name: String, pub(crate) queue_size: usize, } diff --git a/rust/worker/src/memberlist/memberlist_provider.rs b/rust/worker/src/memberlist/memberlist_provider.rs index ea58228ae98..9b529edb8bb 100644 --- a/rust/worker/src/memberlist/memberlist_provider.rs +++ b/rust/worker/src/memberlist/memberlist_provider.rs @@ -1,32 +1,31 @@ -use std::sync::Arc; use std::{fmt::Debug, sync::RwLock}; -use super::config::{CustomResourceMemberlistProviderConfig, MemberlistProviderConfig}; -use crate::system::{Receiver, Sender}; +use super::config::MemberlistProviderConfig; +use crate::system::Receiver; use crate::{ - config::{Configurable, WorkerConfig}, + config::Configurable, errors::{ChromaError, ErrorCodes}, system::{Component, ComponentContext, Handler, StreamHandler}, }; use async_trait::async_trait; -use futures::{StreamExt, TryStreamExt}; -use k8s_openapi::api::events::v1::Event; +use futures::StreamExt; +use kube::runtime::watcher::Config; use kube::{ api::Api, - config, - runtime::{watcher, watcher::Error as WatchError, WatchStreamExt}, + runtime::{watcher, WatchStreamExt}, Client, CustomResource, }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use thiserror::Error; -use tokio_util::sync::CancellationToken; /* =========== Basic Types ============== */ pub(crate) type Memberlist = Vec; #[async_trait] -pub(crate) trait MemberlistProvider: Component + Configurable { +pub(crate) trait MemberlistProvider: + Component + Configurable +{ fn subscribe(&mut self, receiver: Box + Send>) -> (); } @@ -79,7 +78,7 @@ pub(crate) enum CustomResourceMemberlistProviderConfigurationError { impl ChromaError for CustomResourceMemberlistProviderConfigurationError { fn code(&self) -> crate::errors::ErrorCodes { match self { - CustomResourceMemberlistProviderConfigurationError::FailedToLoadKubeClient(e) => { + CustomResourceMemberlistProviderConfigurationError::FailedToLoadKubeClient(_e) => { ErrorCodes::Internal } } @@ -87,9 +86,11 @@ impl ChromaError for CustomResourceMemberlistProviderConfigurationError { } #[async_trait] -impl Configurable for CustomResourceMemberlistProvider { - async fn try_from_config(worker_config: &WorkerConfig) -> Result> { - let my_config = match &worker_config.memberlist_provider { +impl Configurable for CustomResourceMemberlistProvider { + async fn try_from_config( + config: &MemberlistProviderConfig, + ) -> Result> { + let my_config = match &config { MemberlistProviderConfig::CustomResource(config) => config, }; let kube_client = match Client::try_default().await { @@ -102,14 +103,14 @@ impl Configurable for CustomResourceMemberlistProvider { }; let memberlist_cr_client = Api::::namespaced( kube_client.clone(), - &worker_config.kube_namespace, + &my_config.kube_namespace, ); let c: CustomResourceMemberlistProvider = CustomResourceMemberlistProvider { memberlist_name: my_config.memberlist_name.clone(), - kube_ns: worker_config.kube_namespace.clone(), - kube_client: kube_client, - memberlist_cr_client: memberlist_cr_client, + kube_ns: my_config.kube_namespace.clone(), + kube_client, + memberlist_cr_client, queue_size: my_config.queue_size, current_memberlist: RwLock::new(vec![]), subscribers: vec![], @@ -128,11 +129,11 @@ impl CustomResourceMemberlistProvider { let memberlist_cr_client = Api::::namespaced(kube_client.clone(), &kube_ns); CustomResourceMemberlistProvider { - memberlist_name: memberlist_name, - kube_ns: kube_ns, - kube_client: kube_client, - memberlist_cr_client: memberlist_cr_client, - queue_size: queue_size, + memberlist_name, + kube_ns, + kube_client, + memberlist_cr_client, + queue_size, current_memberlist: RwLock::new(vec![]), subscribers: vec![], } @@ -142,7 +143,10 @@ impl CustomResourceMemberlistProvider { let memberlist_cr_client = Api::::namespaced(self.kube_client.clone(), &self.kube_ns); - let stream = watcher(memberlist_cr_client, watcher::Config::default()) + let field_selector = format!("metadata.name={}", self.memberlist_name); + let conifg = Config::default().fields(&field_selector); + + let stream = watcher(memberlist_cr_client, conifg) .default_backoff() .applied_objects(); let stream = stream.then(|event| async move { @@ -164,7 +168,7 @@ impl CustomResourceMemberlistProvider { async fn notify_subscribers(&self) -> () { let curr_memberlist = match self.current_memberlist.read() { Ok(curr_memberlist) => curr_memberlist.clone(), - Err(err) => { + Err(_err) => { // TODO: Log error and attempt recovery return; } @@ -176,12 +180,13 @@ impl CustomResourceMemberlistProvider { } } +#[async_trait] impl Component for CustomResourceMemberlistProvider { fn queue_size(&self) -> usize { self.queue_size } - fn on_start(&mut self, ctx: &ComponentContext) { + async fn on_start(&mut self, ctx: &ComponentContext) { self.connect_to_kube_stream(ctx); } } @@ -217,7 +222,7 @@ impl Handler> for CustomResourceMemberlistProvide Ok(mut curr_memberlist) => { *curr_memberlist = memberlist; } - Err(err) => { + Err(_err) => { // TODO: Log an error } } @@ -257,7 +262,7 @@ mod tests { let kube_ns = "chroma".to_string(); let kube_client = Client::try_default().await.unwrap(); let memberlist_provider = CustomResourceMemberlistProvider::new( - "worker-memberlist".to_string(), + "query-service-memberlist".to_string(), kube_client.clone(), kube_ns.clone(), 10, diff --git a/rust/worker/src/segment/distributed_hnsw_segment.rs b/rust/worker/src/segment/distributed_hnsw_segment.rs index d6f9ca26525..9b8e411839c 100644 --- a/rust/worker/src/segment/distributed_hnsw_segment.rs +++ b/rust/worker/src/segment/distributed_hnsw_segment.rs @@ -1,13 +1,11 @@ -use num_bigint::BigInt; -use parking_lot::{Mutex, RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard}; +use crate::errors::ChromaError; +use crate::index::{HnswIndex, HnswIndexConfig, Index, IndexConfig}; +use crate::types::{LogRecord, Operation, Segment, VectorEmbeddingRecord}; +use parking_lot::RwLock; use std::collections::HashMap; use std::sync::atomic::AtomicUsize; use std::sync::Arc; -use crate::errors::ChromaError; -use crate::index::{HnswIndex, HnswIndexConfig, Index, IndexConfig}; -use crate::types::{EmbeddingRecord, Operation, Segment, VectorEmbeddingRecord}; - pub(crate) struct DistributedHNSWSegment { index: Arc>, id: AtomicUsize, @@ -54,21 +52,21 @@ impl DistributedHNSWSegment { )?)) } - pub(crate) fn write_records(&self, records: Vec>) { - for record in records { - let op = Operation::try_from(record.operation); + pub(crate) fn write_records(&self, log_records: Vec>) { + for log_record in log_records { + let op = Operation::try_from(log_record.record.operation); match op { Ok(Operation::Add) => { // TODO: make lock xor lock - match &record.embedding { + match &log_record.record.embedding { Some(vector) => { let next_id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst); self.user_id_to_id .write() - .insert(record.id.clone(), next_id); + .insert(log_record.record.id.clone(), next_id); self.id_to_user_id .write() - .insert(next_id, record.id.clone()); + .insert(next_id, log_record.record.id.clone()); println!("Segment adding item: {}", next_id); self.index.read().add(next_id, &vector); } @@ -103,11 +101,7 @@ impl DistributedHNSWSegment { let vector = index.get(*internal_id); match vector { Some(vector) => { - let record = VectorEmbeddingRecord { - id: id, - seq_id: BigInt::from(0), - vector, - }; + let record = VectorEmbeddingRecord { id: id, vector }; records.push(Box::new(record)); } None => { diff --git a/rust/worker/src/segment/mod.rs b/rust/worker/src/segment/mod.rs index 5b56f40422f..ebf06b7f31d 100644 --- a/rust/worker/src/segment/mod.rs +++ b/rust/worker/src/segment/mod.rs @@ -1,7 +1,3 @@ pub(crate) mod config; mod distributed_hnsw_segment; -mod segment_ingestor; -mod segment_manager; - -pub(crate) use segment_ingestor::*; -pub(crate) use segment_manager::*; +mod types; diff --git a/rust/worker/src/segment/segment_ingestor.rs b/rust/worker/src/segment/segment_ingestor.rs deleted file mode 100644 index e22abdd9ead..00000000000 --- a/rust/worker/src/segment/segment_ingestor.rs +++ /dev/null @@ -1,48 +0,0 @@ -// A segment ingestor is a component that ingests embeddings into a segment -// Its designed to consume from a async_channel that guarantees exclusive consumption -// They are spawned onto a dedicated thread runtime since ingesting is cpu bound - -use async_trait::async_trait; -use std::{fmt::Debug, sync::Arc}; - -use crate::{ - system::{Component, ComponentContext, ComponentRuntime, Handler}, - types::EmbeddingRecord, -}; - -use super::segment_manager::{self, SegmentManager}; - -pub(crate) struct SegmentIngestor { - segment_manager: SegmentManager, -} - -impl Component for SegmentIngestor { - fn queue_size(&self) -> usize { - 1000 - } - fn runtime() -> ComponentRuntime { - ComponentRuntime::Dedicated - } -} - -impl Debug for SegmentIngestor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("SegmentIngestor").finish() - } -} - -impl SegmentIngestor { - pub(crate) fn new(segment_manager: SegmentManager) -> Self { - SegmentIngestor { - segment_manager: segment_manager, - } - } -} - -#[async_trait] -impl Handler> for SegmentIngestor { - async fn handle(&mut self, message: Box, ctx: &ComponentContext) { - println!("INGEST: ID of embedding is {}", message.id); - self.segment_manager.write_record(message).await; - } -} diff --git a/rust/worker/src/segment/segment_manager.rs b/rust/worker/src/segment/segment_manager.rs deleted file mode 100644 index 314a5b99cdf..00000000000 --- a/rust/worker/src/segment/segment_manager.rs +++ /dev/null @@ -1,252 +0,0 @@ -use crate::{ - config::{Configurable, WorkerConfig}, - errors::ChromaError, - sysdb::sysdb::{GrpcSysDb, SysDb}, - types::VectorQueryResult, -}; -use async_trait::async_trait; -use k8s_openapi::api::node; -use num_bigint::BigInt; -use parking_lot::{ - MappedRwLockReadGuard, RwLock, RwLockReadGuard, RwLockUpgradableReadGuard, RwLockWriteGuard, -}; -use std::collections::HashMap; -use std::sync::Arc; -use uuid::Uuid; - -use super::distributed_hnsw_segment::DistributedHNSWSegment; -use crate::types::{EmbeddingRecord, MetadataValue, Segment, SegmentScope, VectorEmbeddingRecord}; - -#[derive(Clone)] -pub(crate) struct SegmentManager { - inner: Arc, - sysdb: Box, -} - -/// -struct Inner { - vector_segments: RwLock>>, - collection_to_segment_cache: RwLock>>>, - storage_path: Box, -} - -impl SegmentManager { - pub(crate) fn new(sysdb: Box, storage_path: &std::path::Path) -> Self { - SegmentManager { - inner: Arc::new(Inner { - vector_segments: RwLock::new(HashMap::new()), - collection_to_segment_cache: RwLock::new(HashMap::new()), - storage_path: Box::new(storage_path.to_owned()), - }), - sysdb: sysdb, - } - } - - pub(crate) async fn write_record(&mut self, record: Box) { - let collection_id = record.collection_id; - let mut target_segment = None; - // TODO: don't assume 1:1 mapping between collection and segment - { - let segments = self.get_segments(&collection_id).await; - target_segment = match segments { - Ok(found_segments) => { - if found_segments.len() == 0 { - return; // TODO: handle no segment found - } - Some(found_segments[0].clone()) - } - Err(_) => { - // TODO: throw an error and log no segment found - return; - } - }; - } - - let target_segment = match target_segment { - Some(segment) => segment, - None => { - // TODO: throw an error and log no segment found - return; - } - }; - - println!("Writing to segment id {}", target_segment.id); - - let segment_cache = self.inner.vector_segments.upgradable_read(); - match segment_cache.get(&target_segment.id) { - Some(segment) => { - segment.write_records(vec![record]); - } - None => { - let mut segment_cache = RwLockUpgradableReadGuard::upgrade(segment_cache); - - let new_segment = DistributedHNSWSegment::from_segment( - &target_segment, - &self.inner.storage_path, - // TODO: Don't unwrap - throw an error - record.embedding.as_ref().unwrap().len(), - ); - - match new_segment { - Ok(new_segment) => { - new_segment.write_records(vec![record]); - segment_cache.insert(target_segment.id, new_segment); - } - Err(e) => { - println!("Failed to create segment error {}", e); - // TODO: fail and log an error - failed to create/init segment - } - } - } - } - } - - pub(crate) async fn get_records( - &self, - segment_id: &Uuid, - ids: Vec, - ) -> Result>, &'static str> { - // TODO: Load segment if not in cache - let segment_cache = self.inner.vector_segments.read(); - match segment_cache.get(segment_id) { - Some(segment) => { - return Ok(segment.get_records(ids)); - } - None => { - return Err("No segment found"); - } - } - } - - pub(crate) async fn query_vector( - &self, - segment_id: &Uuid, - vectors: &[f32], - k: usize, - include_vector: bool, - ) -> Result>, &'static str> { - let segment_cache = self.inner.vector_segments.read(); - match segment_cache.get(segment_id) { - Some(segment) => { - let mut results = Vec::new(); - let (ids, distances) = segment.query(vectors, k); - for (id, distance) in ids.iter().zip(distances.iter()) { - let fetched_vector = match include_vector { - true => Some(segment.get_records(vec![id.clone()])), - false => None, - }; - - let mut target_record = None; - if include_vector { - target_record = match fetched_vector { - Some(fetched_vectors) => { - if fetched_vectors.len() == 0 { - return Err("No vector found"); - } - let mut target_vec = None; - for vec in fetched_vectors.into_iter() { - if vec.id == *id { - target_vec = Some(vec); - break; - } - } - target_vec - } - None => { - return Err("No vector found"); - } - }; - } - - let ret_vec = match target_record { - Some(target_record) => Some(target_record.vector), - None => None, - }; - - let result = Box::new(VectorQueryResult { - id: id.to_string(), - seq_id: BigInt::from(0), - distance: *distance, - vector: ret_vec, - }); - results.push(result); - } - return Ok(results); - } - None => { - return Err("No segment found"); - } - } - } - - async fn get_segments( - &mut self, - collection_uuid: &Uuid, - ) -> Result>>, &'static str> { - let cache_guard = self.inner.collection_to_segment_cache.read(); - // This lets us return a reference to the segments with the lock. The caller is responsible - // dropping the lock. - let segments = RwLockReadGuard::try_map(cache_guard, |cache| { - return cache.get(&collection_uuid); - }); - match segments { - Ok(segments) => { - return Ok(segments); - } - Err(_) => { - // Data was not in the cache, so we need to get it from the database - // Drop the lock since we need to upgrade it - // Mappable locks cannot be upgraded, so we need to drop the lock and re-acquire it - // https://github.com/Amanieu/parking_lot/issues/83 - drop(segments); - - let segments = self - .sysdb - .get_segments( - None, - None, - Some(SegmentScope::VECTOR), - None, - Some(collection_uuid.clone()), - ) - .await; - match segments { - Ok(segments) => { - let mut cache_guard = self.inner.collection_to_segment_cache.write(); - let mut arc_segments = Vec::new(); - for segment in segments { - arc_segments.push(Arc::new(segment)); - } - cache_guard.insert(collection_uuid.clone(), arc_segments); - let cache_guard = RwLockWriteGuard::downgrade(cache_guard); - let segments = RwLockReadGuard::map(cache_guard, |cache| { - // This unwrap is safe because we just inserted the segments into the cache and currently, - // there is no way to remove segments from the cache. - return cache.get(&collection_uuid).unwrap(); - }); - return Ok(segments); - } - Err(e) => { - return Err("Failed to get segments for collection from SysDB"); - } - } - } - } - } -} - -#[async_trait] -impl Configurable for SegmentManager { - async fn try_from_config(worker_config: &WorkerConfig) -> Result> { - // TODO: Sysdb should have a dynamic resolution in sysdb - let sysdb = GrpcSysDb::try_from_config(worker_config).await; - let sysdb = match sysdb { - Ok(sysdb) => sysdb, - Err(err) => { - return Err(err); - } - }; - let path = std::path::Path::new(&worker_config.segment_manager.storage_path); - Ok(SegmentManager::new(Box::new(sysdb), path)) - } -} diff --git a/rust/worker/src/segment/types.rs b/rust/worker/src/segment/types.rs new file mode 100644 index 00000000000..1f33b1e446d --- /dev/null +++ b/rust/worker/src/segment/types.rs @@ -0,0 +1,12 @@ +use crate::types::LogRecord; + +trait SegmentWriter { + fn begin_transaction(&self); + fn write_records(&self, records: Vec>, offset_ids: Vec); + fn commit_transaction(&self); + fn rollback_transaction(&self); +} + +trait OffsetIdAssigner: SegmentWriter { + fn assign_offset_ids(&self, records: Vec>) -> Vec; +} diff --git a/rust/worker/src/server.rs b/rust/worker/src/server.rs index 1ecc6ba2e70..2ed74efb4e6 100644 --- a/rust/worker/src/server.rs +++ b/rust/worker/src/server.rs @@ -1,28 +1,52 @@ -use std::f32::consts::E; - use crate::chroma_proto; use crate::chroma_proto::{ GetVectorsRequest, GetVectorsResponse, QueryVectorsRequest, QueryVectorsResponse, }; -use crate::config::{Configurable, WorkerConfig}; +use crate::config::{Configurable, QueryServiceConfig}; use crate::errors::ChromaError; -use crate::segment::SegmentManager; +use crate::execution::operator::TaskMessage; +use crate::execution::orchestration::HnswQueryOrchestrator; +use crate::log::log::Log; +use crate::sysdb::sysdb::SysDb; +use crate::system::{Receiver, System}; use crate::types::ScalarEncoding; use async_trait::async_trait; -use kube::core::request; use tonic::{transport::Server, Request, Response, Status}; use uuid::Uuid; pub struct WorkerServer { - segment_manager: Option, + // System + system: Option, + // Component dependencies + dispatcher: Option>>, + // Service dependencies + log: Box, + sysdb: Box, port: u16, } #[async_trait] -impl Configurable for WorkerServer { - async fn try_from_config(config: &WorkerConfig) -> Result> { +impl Configurable for WorkerServer { + async fn try_from_config(config: &QueryServiceConfig) -> Result> { + let sysdb_config = &config.sysdb; + let sysdb = match crate::sysdb::from_config(sysdb_config).await { + Ok(sysdb) => sysdb, + Err(err) => { + return Err(err); + } + }; + let log_config = &config.log; + let log = match crate::log::from_config(log_config).await { + Ok(log) => log, + Err(err) => { + return Err(err); + } + }; Ok(WorkerServer { - segment_manager: None, + dispatcher: None, + system: None, + sysdb, + log, port: config.my_port, }) } @@ -32,7 +56,7 @@ impl WorkerServer { pub(crate) async fn run(worker: WorkerServer) -> Result<(), Box> { let addr = format!("[::]:{}", worker.port).parse().unwrap(); println!("Worker listening on {}", addr); - let server = Server::builder() + let _server = Server::builder() .add_service(chroma_proto::vector_reader_server::VectorReaderServer::new( worker, )) @@ -43,8 +67,12 @@ impl WorkerServer { Ok(()) } - pub(crate) fn set_segment_manager(&mut self, segment_manager: SegmentManager) { - self.segment_manager = Some(segment_manager); + pub(crate) fn set_dispatcher(&mut self, dispatcher: Box>) { + self.dispatcher = Some(dispatcher); + } + + pub(crate) fn set_system(&mut self, system: System) { + self.system = Some(system); } } @@ -55,55 +83,14 @@ impl chroma_proto::vector_reader_server::VectorReader for WorkerServer { request: Request, ) -> Result, Status> { let request = request.into_inner(); - let segment_uuid = match Uuid::parse_str(&request.segment_id) { + let _segment_uuid = match Uuid::parse_str(&request.segment_id) { Ok(uuid) => uuid, Err(_) => { return Err(Status::invalid_argument("Invalid UUID")); } }; - let segment_manager = match self.segment_manager { - Some(ref segment_manager) => segment_manager, - None => { - return Err(Status::internal("No segment manager found")); - } - }; - - let records = match segment_manager - .get_records(&segment_uuid, request.ids) - .await - { - Ok(records) => records, - Err(e) => { - return Err(Status::internal(format!("Error getting records: {}", e))); - } - }; - - let mut proto_records = Vec::new(); - for record in records { - let sed_id_bytes = record.seq_id.to_bytes_le(); - let dim = record.vector.len(); - let proto_vector = (record.vector, ScalarEncoding::FLOAT32, dim).try_into(); - match proto_vector { - Ok(proto_vector) => { - let proto_record = chroma_proto::VectorEmbeddingRecord { - id: record.id, - seq_id: sed_id_bytes.1, - vector: Some(proto_vector), - }; - proto_records.push(proto_record); - } - Err(e) => { - return Err(Status::internal(format!("Error converting vector: {}", e))); - } - } - } - - let resp = chroma_proto::GetVectorsResponse { - records: proto_records, - }; - - Ok(Response::new(resp)) + Err(Status::unimplemented("Not yet implemented")) } async fn query_vectors( @@ -118,46 +105,67 @@ impl chroma_proto::vector_reader_server::VectorReader for WorkerServer { } }; - let segment_manager = match self.segment_manager { - Some(ref segment_manager) => segment_manager, - None => { - return Err(Status::internal("No segment manager found")); - } - }; - let mut proto_results_for_all = Vec::new(); + + let mut query_vectors = Vec::new(); for proto_query_vector in request.vectors { - let (query_vector, encoding) = match proto_query_vector.try_into() { + let (query_vector, _encoding) = match proto_query_vector.try_into() { Ok((vector, encoding)) => (vector, encoding), Err(e) => { return Err(Status::internal(format!("Error converting vector: {}", e))); } }; + query_vectors.push(query_vector); + } + + let dispatcher = match self.dispatcher { + Some(ref dispatcher) => dispatcher, + None => { + return Err(Status::internal("No dispatcher found")); + } + }; - let results = match segment_manager - .query_vector( - &segment_uuid, - &query_vector, - request.k as usize, + let result = match self.system { + Some(ref system) => { + let orchestrator = HnswQueryOrchestrator::new( + // TODO: Should not have to clone query vectors here + system.clone(), + query_vectors.clone(), + request.k, request.include_embeddings, - ) - .await - { - Ok(results) => results, - Err(e) => { - return Err(Status::internal(format!("Error querying segment: {}", e))); - } - }; + segment_uuid, + self.log.clone(), + self.sysdb.clone(), + dispatcher.clone(), + ); + orchestrator.run().await + } + None => { + return Err(Status::internal("No system found")); + } + }; + let result = match result { + Ok(result) => result, + Err(e) => { + return Err(Status::internal(format!( + "Error running orchestrator: {}", + e + ))); + } + }; + + for result_set in result { let mut proto_results = Vec::new(); - for query_result in results { + for query_result in result_set { let proto_result = chroma_proto::VectorQueryResult { id: query_result.id, - seq_id: query_result.seq_id.to_bytes_le().1, distance: query_result.distance, vector: match query_result.vector { Some(vector) => { - match (vector, ScalarEncoding::FLOAT32, query_vector.len()).try_into() { + match (vector, ScalarEncoding::FLOAT32, query_vectors[0].len()) + .try_into() + { Ok(proto_vector) => Some(proto_vector), Err(e) => { return Err(Status::internal(format!( @@ -172,11 +180,9 @@ impl chroma_proto::vector_reader_server::VectorReader for WorkerServer { }; proto_results.push(proto_result); } - - let vector_query_results = chroma_proto::VectorQueryResults { + proto_results_for_all.push(chroma_proto::VectorQueryResults { results: proto_results, - }; - proto_results_for_all.push(vector_query_results); + }); } let resp = chroma_proto::QueryVectorsResponse { diff --git a/rust/worker/src/storage/s3.rs b/rust/worker/src/storage/s3.rs index f78767e4896..0a43388d491 100644 --- a/rust/worker/src/storage/s3.rs +++ b/rust/worker/src/storage/s3.rs @@ -9,7 +9,7 @@ // streaming from s3. use super::{config::StorageConfig, Storage}; -use crate::config::{Configurable, WorkerConfig}; +use crate::config::Configurable; use crate::errors::ChromaError; use async_trait::async_trait; use aws_sdk_s3; @@ -73,9 +73,9 @@ impl S3Storage { } #[async_trait] -impl Configurable for S3Storage { - async fn try_from_config(config: &WorkerConfig) -> Result> { - match &config.storage { +impl Configurable for S3Storage { + async fn try_from_config(config: &StorageConfig) -> Result> { + match &config { StorageConfig::S3(s3_config) => { let config = aws_config::load_from_env().await; let client = aws_sdk_s3::Client::new(&config); @@ -90,7 +90,7 @@ impl Configurable for S3Storage { #[async_trait] impl Storage for S3Storage { async fn get(&self, key: &str, path: &str) -> Result<(), String> { - let mut file = std::fs::File::create(path); + let file = std::fs::File::create(path); let res = self .client .get_object() diff --git a/rust/worker/src/sysdb/mod.rs b/rust/worker/src/sysdb/mod.rs index 1db5510f893..21e22265568 100644 --- a/rust/worker/src/sysdb/mod.rs +++ b/rust/worker/src/sysdb/mod.rs @@ -1,2 +1,17 @@ pub(crate) mod config; pub(crate) mod sysdb; +pub(crate) mod test_sysdb; + +use self::config::SysDbConfig; +use crate::config::Configurable; +use crate::errors::ChromaError; + +pub(crate) async fn from_config( + config: &SysDbConfig, +) -> Result, Box> { + match &config { + crate::sysdb::config::SysDbConfig::Grpc(_) => { + Ok(Box::new(sysdb::GrpcSysDb::try_from_config(config).await?)) + } + } +} diff --git a/rust/worker/src/sysdb/sysdb.rs b/rust/worker/src/sysdb/sysdb.rs index ba8be18fdf5..61a802a006e 100644 --- a/rust/worker/src/sysdb/sysdb.rs +++ b/rust/worker/src/sysdb/sysdb.rs @@ -1,27 +1,25 @@ -use async_trait::async_trait; -use uuid::Uuid; - +use super::config::SysDbConfig; use crate::chroma_proto; -use crate::config::{Configurable, WorkerConfig}; +use crate::config::Configurable; use crate::types::{CollectionConversionError, SegmentConversionError}; use crate::{ chroma_proto::sys_db_client, errors::{ChromaError, ErrorCodes}, types::{Collection, Segment, SegmentScope}, }; +use async_trait::async_trait; +use std::fmt::Debug; use thiserror::Error; - -use super::config::SysDbConfig; +use uuid::Uuid; const DEFAULT_DATBASE: &str = "default_database"; const DEFAULT_TENANT: &str = "default_tenant"; #[async_trait] -pub(crate) trait SysDb: Send + Sync + SysDbClone { +pub(crate) trait SysDb: Send + Sync + SysDbClone + Debug { async fn get_collections( &mut self, collection_id: Option, - topic: Option, name: Option, tenant: Option, database: Option, @@ -32,7 +30,6 @@ pub(crate) trait SysDb: Send + Sync + SysDbClone { id: Option, r#type: Option, scope: Option, - topic: Option, collection: Option, ) -> Result, GetSegmentsError>; } @@ -59,7 +56,7 @@ impl Clone for Box { } } -#[derive(Clone)] +#[derive(Clone, Debug)] // Since this uses tonic transport channel, cloning is cheap. Each client only supports // one inflight request at a time, so we need to clone the client for each requester. pub(crate) struct GrpcSysDb { @@ -81,9 +78,9 @@ impl ChromaError for GrpcSysDbError { } #[async_trait] -impl Configurable for GrpcSysDb { - async fn try_from_config(worker_config: &WorkerConfig) -> Result> { - match &worker_config.sysdb { +impl Configurable for GrpcSysDb { + async fn try_from_config(config: &SysDbConfig) -> Result> { + match &config { SysDbConfig::Grpc(my_config) => { let host = &my_config.host; let port = &my_config.port; @@ -108,7 +105,6 @@ impl SysDb for GrpcSysDb { async fn get_collections( &mut self, collection_id: Option, - topic: Option, name: Option, tenant: Option, database: Option, @@ -128,7 +124,6 @@ impl SysDb for GrpcSysDb { .client .get_collections(chroma_proto::GetCollectionsRequest { id: collection_id_str, - topic: topic, name: name, tenant: if tenant.is_some() { tenant.unwrap() @@ -172,7 +167,6 @@ impl SysDb for GrpcSysDb { id: Option, r#type: Option, scope: Option, - topic: Option, collection: Option, ) -> Result, GetSegmentsError> { let res = self @@ -190,7 +184,6 @@ impl SysDb for GrpcSysDb { } else { None }, - topic: topic, collection: if collection.is_some() { Some(collection.unwrap().to_string()) } else { diff --git a/rust/worker/src/sysdb/test_sysdb.rs b/rust/worker/src/sysdb/test_sysdb.rs new file mode 100644 index 00000000000..9f00f5b42cb --- /dev/null +++ b/rust/worker/src/sysdb/test_sysdb.rs @@ -0,0 +1,84 @@ +use crate::sysdb::sysdb::GetCollectionsError; +use crate::sysdb::sysdb::GetSegmentsError; +use crate::sysdb::sysdb::SysDb; +use crate::types::Collection; +use crate::types::Segment; +use crate::types::SegmentScope; +use async_trait::async_trait; +use std::collections::HashMap; +use uuid::Uuid; + +#[derive(Clone, Debug)] +pub(crate) struct TestSysDb { + collections: HashMap, +} + +impl TestSysDb { + pub(crate) fn new() -> Self { + TestSysDb { + collections: HashMap::new(), + } + } + + pub(crate) fn add_collection(&mut self, collection: Collection) { + self.collections.insert(collection.id, collection); + } + + fn filter_collections( + collection: &Collection, + collection_id: Option, + name: Option, + tenant: Option, + database: Option, + ) -> bool { + if collection_id.is_some() && collection_id.unwrap() != collection.id { + return false; + } + if name.is_some() && name.unwrap() != collection.name { + return false; + } + if tenant.is_some() && tenant.unwrap() != collection.tenant { + return false; + } + if database.is_some() && database.unwrap() != collection.database { + return false; + } + true + } +} + +#[async_trait] +impl SysDb for TestSysDb { + async fn get_collections( + &mut self, + collection_id: Option, + name: Option, + tenant: Option, + database: Option, + ) -> Result, GetCollectionsError> { + let mut collections = Vec::new(); + for collection in self.collections.values() { + if !TestSysDb::filter_collections( + &collection, + collection_id, + name.clone(), + tenant.clone(), + database.clone(), + ) { + continue; + } + collections.push(collection.clone()); + } + Ok(collections) + } + + async fn get_segments( + &mut self, + _id: Option, + _type: Option, + _scope: Option, + _collection: Option, + ) -> Result, GetSegmentsError> { + Ok(Vec::new()) + } +} diff --git a/rust/worker/src/system/executor.rs b/rust/worker/src/system/executor.rs index c50ac4a56fa..4877273b70e 100644 --- a/rust/worker/src/system/executor.rs +++ b/rust/worker/src/system/executor.rs @@ -1,14 +1,12 @@ -use std::sync::Arc; - -use tokio::select; - -use crate::system::ComponentContext; - use super::{ + scheduler::Scheduler, sender::{Sender, Wrapper}, system::System, Component, }; +use crate::system::ComponentContext; +use std::sync::Arc; +use tokio::select; struct Inner where @@ -17,6 +15,7 @@ where pub(super) sender: Sender, pub(super) cancellation_token: tokio_util::sync::CancellationToken, pub(super) system: System, + pub(super) scheduler: Scheduler, } #[derive(Clone)] @@ -40,38 +39,49 @@ where cancellation_token: tokio_util::sync::CancellationToken, handler: C, system: System, + scheduler: Scheduler, ) -> Self { ComponentExecutor { inner: Arc::new(Inner { sender, cancellation_token, system, + scheduler, }), handler, } } pub(super) async fn run(&mut self, mut channel: tokio::sync::mpsc::Receiver>) { + self.handler + .on_start(&ComponentContext { + system: self.inner.system.clone(), + sender: self.inner.sender.clone(), + cancellation_token: self.inner.cancellation_token.clone(), + scheduler: self.inner.scheduler.clone(), + }) + .await; loop { select! { - _ = self.inner.cancellation_token.cancelled() => { - break; - } - message = channel.recv() => { - match message { - Some(mut message) => { - message.handle(&mut self.handler, - &ComponentContext{ - system: self.inner.system.clone(), - sender: self.inner.sender.clone(), - cancellation_token: self.inner.cancellation_token.clone(), - } - ).await; - } - None => { - // TODO: Log error - } + _ = self.inner.cancellation_token.cancelled() => { + break; + } + message = channel.recv() => { + match message { + Some(mut message) => { + message.handle(&mut self.handler, + &ComponentContext{ + system: self.inner.system.clone(), + sender: self.inner.sender.clone(), + cancellation_token: self.inner.cancellation_token.clone(), + scheduler: self.inner.scheduler.clone(), + } + ).await; + } + None => { + // TODO: Log error } + } } } } diff --git a/rust/worker/src/system/mod.rs b/rust/worker/src/system/mod.rs index 32ad862f768..9656b0291c2 100644 --- a/rust/worker/src/system/mod.rs +++ b/rust/worker/src/system/mod.rs @@ -1,4 +1,5 @@ mod executor; +mod scheduler; mod sender; mod system; mod types; diff --git a/rust/worker/src/system/scheduler.rs b/rust/worker/src/system/scheduler.rs new file mode 100644 index 00000000000..34cb3b1872a --- /dev/null +++ b/rust/worker/src/system/scheduler.rs @@ -0,0 +1,211 @@ +use parking_lot::RwLock; +use std::fmt::Debug; +use std::sync::Arc; +use std::time::Duration; +use tokio::select; + +use super::sender::Sender; +use super::{Component, ComponentContext, Handler}; + +#[derive(Debug)] +pub(crate) struct SchedulerTaskHandle { + join_handle: Option>, + cancel: tokio_util::sync::CancellationToken, +} + +#[derive(Clone, Debug)] +pub(crate) struct Scheduler { + handles: Arc>>, +} + +impl Scheduler { + pub(crate) fn new() -> Scheduler { + Scheduler { + handles: Arc::new(RwLock::new(Vec::new())), + } + } + + pub(crate) fn schedule( + &self, + sender: Sender, + message: M, + duration: Duration, + ctx: &ComponentContext, + ) where + C: Component + Handler, + M: Debug + Send + 'static, + { + let cancel = ctx.cancellation_token.clone(); + let handle = tokio::spawn(async move { + select! { + _ = cancel.cancelled() => { + return; + } + _ = tokio::time::sleep(duration) => { + match sender.send(message).await { + Ok(_) => { + return; + }, + Err(e) => { + // TODO: log error + println!("Error: {:?}", e); + return; + } + } + } + } + }); + let handle = SchedulerTaskHandle { + join_handle: Some(handle), + cancel: ctx.cancellation_token.clone(), + }; + self.handles.write().push(handle); + } + + pub(crate) fn schedule_interval( + &self, + sender: Sender, + message: M, + duration: Duration, + num_times: Option, + ctx: &ComponentContext, + ) where + C: Component + Handler, + M: Debug + Send + Clone + 'static, + { + let cancel = ctx.cancellation_token.clone(); + + let handle = tokio::spawn(async move { + let mut counter = 0; + while Self::should_continue(num_times, counter) { + select! { + _ = cancel.cancelled() => { + return; + } + _ = tokio::time::sleep(duration) => { + match sender.send(message.clone()).await { + Ok(_) => { + }, + Err(e) => { + // TODO: log error + println!("Error: {:?}", e); + } + } + } + } + counter += 1; + } + }); + let handle = SchedulerTaskHandle { + join_handle: Some(handle), + cancel: ctx.cancellation_token.clone(), + }; + self.handles.write().push(handle); + } + + fn should_continue(num_times: Option, counter: usize) -> bool { + if num_times.is_some() { + let num_times = num_times.unwrap(); + if counter >= num_times { + return false; + } + } + true + } + + // Note: this method holds the lock on the handles, should call it only after stop is + // called. + pub(crate) async fn join(&self) { + let mut handles = self.handles.write(); + for handle in handles.iter_mut() { + if let Some(join_handle) = handle.join_handle.take() { + match join_handle.await { + Ok(_) => {} + Err(e) => { + println!("Error: {:?}", e); + } + } + } + } + } + + pub(crate) fn stop(&self) { + let handles = self.handles.read(); + for handle in handles.iter() { + handle.cancel.cancel(); + } + } +} + +mod tests { + use super::*; + use crate::system::System; + use async_trait::async_trait; + use std::sync::Arc; + use std::time::Duration; + + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[derive(Debug)] + struct TestComponent { + queue_size: usize, + counter: Arc, + } + + #[derive(Clone, Debug)] + struct ScheduleMessage {} + + impl TestComponent { + fn new(queue_size: usize, counter: Arc) -> Self { + TestComponent { + queue_size, + counter, + } + } + } + #[async_trait] + impl Handler for TestComponent { + async fn handle( + &mut self, + _message: ScheduleMessage, + _ctx: &ComponentContext, + ) { + self.counter.fetch_add(1, Ordering::SeqCst); + } + } + + #[async_trait] + impl Component for TestComponent { + fn queue_size(&self) -> usize { + self.queue_size + } + + async fn on_start(&mut self, ctx: &ComponentContext) -> () { + let duration = Duration::from_millis(100); + ctx.scheduler + .schedule(ctx.sender.clone(), ScheduleMessage {}, duration, ctx); + + let num_times = 4; + ctx.scheduler.schedule_interval( + ctx.sender.clone(), + ScheduleMessage {}, + duration, + Some(num_times), + ctx, + ); + } + } + + #[tokio::test] + async fn test_schedule() { + let system = System::new(); + let counter = Arc::new(AtomicUsize::new(0)); + let component = TestComponent::new(10, counter.clone()); + let _handle = system.start_component(component); + // yield to allow the component to process the messages + tokio::task::yield_now().await; + // We should have scheduled the message once + system.join().await; + assert_eq!(counter.load(Ordering::SeqCst), 5); + } +} diff --git a/rust/worker/src/system/sender.rs b/rust/worker/src/system/sender.rs index df2e1bc5587..d9fb0785418 100644 --- a/rust/worker/src/system/sender.rs +++ b/rust/worker/src/system/sender.rs @@ -51,7 +51,6 @@ where } // Sender - pub(crate) struct Sender where C: Component + Send + 'static, @@ -78,6 +77,14 @@ where Err(_) => Err(ChannelError::SendError), } } + + pub(crate) fn as_receiver(&self) -> Box> + where + C: Component + Handler, + M: Debug + Send + 'static, + { + Box::new(ReceiverImpl::new(self.sender.clone())) + } } impl Clone for Sender @@ -94,7 +101,7 @@ where // Reciever Traits #[async_trait] -pub(crate) trait Receiver: Send + Sync + ReceiverClone { +pub(crate) trait Receiver: Send + Sync + Debug + ReceiverClone { async fn send(&self, message: M) -> Result<(), ChannelError>; } @@ -118,7 +125,7 @@ where } // Reciever Impls - +#[derive(Debug)] pub(super) struct ReceiverImpl where C: Component, diff --git a/rust/worker/src/system/system.rs b/rust/worker/src/system/system.rs index 238da52a0ee..0d9f4738625 100644 --- a/rust/worker/src/system/system.rs +++ b/rust/worker/src/system/system.rs @@ -1,60 +1,58 @@ -use std::fmt::Debug; -use std::sync::Arc; - +use super::scheduler::Scheduler; +use super::sender::Sender; +use super::ComponentContext; +use super::ComponentRuntime; +use super::{executor::ComponentExecutor, Component, ComponentHandle, Handler, StreamHandler}; use futures::Stream; use futures::StreamExt; +use std::fmt::Debug; +use std::sync::Arc; use tokio::runtime::Builder; use tokio::{pin, select}; -use super::ComponentRuntime; -// use super::executor::StreamComponentExecutor; -use super::sender::{self, Sender, Wrapper}; -use super::{executor, ComponentContext}; -use super::{executor::ComponentExecutor, Component, ComponentHandle, Handler, StreamHandler}; -use std::sync::Mutex; - -#[derive(Clone)] +#[derive(Clone, Debug)] pub(crate) struct System { - inner: Arc>, + inner: Arc, } -struct Inner {} +#[derive(Debug)] +struct Inner { + scheduler: Scheduler, +} impl System { pub(crate) fn new() -> System { System { - inner: Arc::new(Mutex::new(Inner {})), + inner: Arc::new(Inner { + scheduler: Scheduler::new(), + }), } } - pub(crate) fn start_component(&mut self, mut component: C) -> ComponentHandle + pub(crate) fn start_component(&self, component: C) -> ComponentHandle where C: Component + Send + 'static, { let (tx, rx) = tokio::sync::mpsc::channel(component.queue_size()); let sender = Sender::new(tx); let cancel_token = tokio_util::sync::CancellationToken::new(); - let _ = component.on_start(&ComponentContext { - system: self.clone(), - sender: sender.clone(), - cancellation_token: cancel_token.clone(), - }); let mut executor = ComponentExecutor::new( sender.clone(), cancel_token.clone(), component, self.clone(), + self.inner.scheduler.clone(), ); match C::runtime() { - ComponentRuntime::Global => { + ComponentRuntime::Inherit => { let join_handle = tokio::spawn(async move { executor.run(rx).await }); return ComponentHandle::new(cancel_token, Some(join_handle), sender); } ComponentRuntime::Dedicated => { println!("Spawning on dedicated thread"); // Spawn on a dedicated thread - let mut rt = Builder::new_current_thread().enable_all().build().unwrap(); + let rt = Builder::new_current_thread().enable_all().build().unwrap(); let join_handle = std::thread::spawn(move || { rt.block_on(async move { executor.run(rx).await }); }); @@ -74,9 +72,18 @@ impl System { system: self.clone(), sender: ctx.sender.clone(), cancellation_token: ctx.cancellation_token.clone(), + scheduler: ctx.scheduler.clone(), }; tokio::spawn(async move { stream_loop(stream, &ctx).await }); } + + pub(crate) async fn stop(&self) { + self.inner.scheduler.stop(); + } + + pub(crate) async fn join(&self) { + self.inner.scheduler.join().await; + } } async fn stream_loop(stream: S, ctx: &ComponentContext) diff --git a/rust/worker/src/system/types.rs b/rust/worker/src/system/types.rs index 9c2cd463561..e1e2228b146 100644 --- a/rust/worker/src/system/types.rs +++ b/rust/worker/src/system/types.rs @@ -1,12 +1,9 @@ -use std::{fmt::Debug, sync::Arc}; - +use super::scheduler::Scheduler; use async_trait::async_trait; use futures::Stream; -use tokio::select; +use std::fmt::Debug; -use super::{ - executor::ComponentExecutor, sender::Sender, system::System, Receiver, ReceiverImpl, Wrapper, -}; +use super::{sender::Sender, system::System, Receiver, ReceiverImpl}; #[derive(Debug, PartialEq)] /// The state of a component @@ -20,7 +17,7 @@ pub(crate) enum ComponentState { #[derive(Debug, PartialEq)] pub(crate) enum ComponentRuntime { - Global, + Inherit, Dedicated, } @@ -33,12 +30,13 @@ pub(crate) enum ComponentRuntime { /// # Methods /// - queue_size: The size of the queue to use for the component before it starts dropping messages /// - on_start: Called when the component is started +#[async_trait] pub(crate) trait Component: Send + Sized + Debug + 'static { fn queue_size(&self) -> usize; fn runtime() -> ComponentRuntime { - ComponentRuntime::Global + ComponentRuntime::Inherit } - fn on_start(&mut self, ctx: &ComponentContext) -> () {} + async fn on_start(&mut self, _ctx: &ComponentContext) -> () {} } /// A handler is a component that can process messages of a given type. @@ -135,6 +133,7 @@ where pub(crate) system: System, pub(crate) sender: Sender, pub(crate) cancellation_token: tokio_util::sync::CancellationToken, + pub(crate) scheduler: Scheduler, } #[cfg(test)] @@ -142,8 +141,8 @@ mod tests { use super::*; use async_trait::async_trait; use futures::stream; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; #[derive(Debug)] struct TestComponent { @@ -154,8 +153,8 @@ mod tests { impl TestComponent { fn new(queue_size: usize, counter: Arc) -> Self { TestComponent { - queue_size: queue_size, - counter: counter, + queue_size, + counter, } } } @@ -166,15 +165,15 @@ mod tests { self.counter.fetch_add(message, Ordering::SeqCst); } } - impl StreamHandler for TestComponent {} + #[async_trait] impl Component for TestComponent { fn queue_size(&self) -> usize { - return self.queue_size; + self.queue_size } - fn on_start(&mut self, ctx: &ComponentContext) -> () { + async fn on_start(&mut self, ctx: &ComponentContext) -> () { let test_stream = stream::iter(vec![1, 2, 3]); self.register_stream(test_stream, ctx); } @@ -182,7 +181,7 @@ mod tests { #[tokio::test] async fn it_can_work() { - let mut system = System::new(); + let system = System::new(); let counter = Arc::new(AtomicUsize::new(0)); let component = TestComponent::new(10, counter.clone()); let mut handle = system.start_component(component); @@ -191,12 +190,13 @@ mod tests { handle.sender.send(3).await.unwrap(); // yield to allow the component to process the messages tokio::task::yield_now().await; + // With the streaming data and the messages we should have 12 + assert_eq!(counter.load(Ordering::SeqCst), 12); handle.stop(); // Yield to allow the component to stop tokio::task::yield_now().await; + // Expect the component to be stopped assert_eq!(*handle.state(), ComponentState::Stopped); - // With the streaming data and the messages we should have 12 - assert_eq!(counter.load(Ordering::SeqCst), 12); let res = handle.sender.send(4).await; // Expect an error because the component is stopped assert!(res.is_err()); diff --git a/rust/worker/src/types/collection.rs b/rust/worker/src/types/collection.rs index 2dd495a5afc..1f8766a1ee1 100644 --- a/rust/worker/src/types/collection.rs +++ b/rust/worker/src/types/collection.rs @@ -6,15 +6,16 @@ use crate::{ use thiserror::Error; use uuid::Uuid; -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) struct Collection { pub(crate) id: Uuid, pub(crate) name: String, - pub(crate) topic: String, pub(crate) metadata: Option, pub(crate) dimension: Option, pub(crate) tenant: String, pub(crate) database: String, + pub(crate) log_position: i64, + pub(crate) version: i32, } #[derive(Error, Debug)] @@ -52,11 +53,12 @@ impl TryFrom for Collection { Ok(Collection { id: collection_uuid, name: proto_collection.name, - topic: proto_collection.topic, metadata: collection_metadata, dimension: proto_collection.dimension, tenant: proto_collection.tenant, database: proto_collection.database, + log_position: proto_collection.log_position, + version: proto_collection.version, }) } } @@ -70,16 +72,16 @@ mod test { let proto_collection = chroma_proto::Collection { id: "00000000-0000-0000-0000-000000000000".to_string(), name: "foo".to_string(), - topic: "bar".to_string(), metadata: None, dimension: None, tenant: "baz".to_string(), database: "qux".to_string(), + log_position: 0, + version: 0, }; let converted_collection: Collection = proto_collection.try_into().unwrap(); assert_eq!(converted_collection.id, Uuid::nil()); assert_eq!(converted_collection.name, "foo".to_string()); - assert_eq!(converted_collection.topic, "bar".to_string()); assert_eq!(converted_collection.metadata, None); assert_eq!(converted_collection.dimension, None); assert_eq!(converted_collection.tenant, "baz".to_string()); diff --git a/rust/worker/src/types/metadata.rs b/rust/worker/src/types/metadata.rs index 73f4c749e1e..f792a55ad4b 100644 --- a/rust/worker/src/types/metadata.rs +++ b/rust/worker/src/types/metadata.rs @@ -5,7 +5,7 @@ use crate::{ use std::collections::HashMap; use thiserror::Error; -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) enum UpdateMetadataValue { Int(i32), Float(f64), @@ -52,7 +52,7 @@ MetadataValue =========================================== */ -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) enum MetadataValue { Int(i32), Float(f64), @@ -130,7 +130,6 @@ impl TryFrom<&chroma_proto::UpdateMetadataValue> for MetadataValue { UpdateMetadata =========================================== */ - pub(crate) type UpdateMetadata = HashMap; impl TryFrom for UpdateMetadata { diff --git a/rust/worker/src/types/mod.rs b/rust/worker/src/types/mod.rs index edda924c42c..b31be7792d0 100644 --- a/rust/worker/src/types/mod.rs +++ b/rust/worker/src/types/mod.rs @@ -1,19 +1,19 @@ #[macro_use] mod types; mod collection; -mod embedding_record; mod metadata; mod operation; +mod record; mod scalar_encoding; mod segment; mod segment_scope; // Re-export the types module, so that we can use it as a single import in other modules. -pub use collection::*; -pub use embedding_record::*; -pub use metadata::*; -pub use operation::*; -pub use scalar_encoding::*; -pub use segment::*; -pub use segment_scope::*; -pub use types::*; +pub(crate) use collection::*; +pub(crate) use metadata::*; +pub(crate) use operation::*; +pub(crate) use record::*; +pub(crate) use scalar_encoding::*; +pub(crate) use segment::*; +pub(crate) use segment_scope::*; +pub(crate) use types::*; diff --git a/rust/worker/src/types/operation.rs b/rust/worker/src/types/operation.rs index 581e5c39f8e..b0bec9cf9d5 100644 --- a/rust/worker/src/types/operation.rs +++ b/rust/worker/src/types/operation.rs @@ -5,7 +5,7 @@ use crate::{ }; use thiserror::Error; -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) enum Operation { Add, Update, diff --git a/rust/worker/src/types/embedding_record.rs b/rust/worker/src/types/record.rs similarity index 59% rename from rust/worker/src/types/embedding_record.rs rename to rust/worker/src/types/record.rs index 14957a85349..86b6c912fd0 100644 --- a/rust/worker/src/types/embedding_record.rs +++ b/rust/worker/src/types/record.rs @@ -1,29 +1,30 @@ use super::{ ConversionError, Operation, OperationConversionError, ScalarEncoding, - ScalarEncodingConversionError, SeqId, UpdateMetadata, UpdateMetadataValueConversionError, + ScalarEncodingConversionError, UpdateMetadata, UpdateMetadataValueConversionError, }; use crate::{ chroma_proto, errors::{ChromaError, ErrorCodes}, }; use thiserror::Error; -use uuid::Uuid; -#[derive(Debug)] -pub(crate) struct EmbeddingRecord { +#[derive(Clone, Debug)] +pub(crate) struct OperationRecord { pub(crate) id: String, - pub(crate) seq_id: SeqId, - pub(crate) embedding: Option>, // NOTE: we only support float32 embeddings for now + pub(crate) embedding: Option>, // NOTE: we only support float32 embeddings for now so this ignores the encoding pub(crate) encoding: Option, pub(crate) metadata: Option, pub(crate) operation: Operation, - pub(crate) collection_id: Uuid, } -pub(crate) type SubmitEmbeddingRecordWithSeqId = (chroma_proto::SubmitEmbeddingRecord, SeqId); +#[derive(Clone, Debug)] +pub(crate) struct LogRecord { + pub(crate) log_offset: i64, + pub(crate) record: OperationRecord, +} #[derive(Error, Debug)] -pub(crate) enum EmbeddingRecordConversionError { +pub(crate) enum RecordConversionError { #[error("Invalid UUID")] InvalidUuid, #[error(transparent)] @@ -38,61 +39,67 @@ pub(crate) enum EmbeddingRecordConversionError { VectorConversionError(#[from] VectorConversionError), } -impl_base_convert_error!(EmbeddingRecordConversionError, { - EmbeddingRecordConversionError::InvalidUuid => ErrorCodes::InvalidArgument, - EmbeddingRecordConversionError::OperationConversionError(inner) => inner.code(), - EmbeddingRecordConversionError::ScalarEncodingConversionError(inner) => inner.code(), - EmbeddingRecordConversionError::UpdateMetadataValueConversionError(inner) => inner.code(), - EmbeddingRecordConversionError::VectorConversionError(inner) => inner.code(), +impl_base_convert_error!(RecordConversionError, { + RecordConversionError::InvalidUuid => ErrorCodes::InvalidArgument, + RecordConversionError::OperationConversionError(inner) => inner.code(), + RecordConversionError::ScalarEncodingConversionError(inner) => inner.code(), + RecordConversionError::UpdateMetadataValueConversionError(inner) => inner.code(), + RecordConversionError::VectorConversionError(inner) => inner.code(), }); -impl TryFrom for EmbeddingRecord { - type Error = EmbeddingRecordConversionError; +impl TryFrom for OperationRecord { + type Error = RecordConversionError; fn try_from( - proto_submit_with_seq_id: SubmitEmbeddingRecordWithSeqId, + operation_record_proto: chroma_proto::OperationRecord, ) -> Result { - let proto_submit = proto_submit_with_seq_id.0; - let seq_id = proto_submit_with_seq_id.1; - let op = match proto_submit.operation.try_into() { + let operation = match operation_record_proto.operation.try_into() { Ok(op) => op, - Err(e) => return Err(EmbeddingRecordConversionError::OperationConversionError(e)), - }; - - let collection_uuid = match Uuid::try_parse(&proto_submit.collection_id) { - Ok(uuid) => uuid, - Err(_) => return Err(EmbeddingRecordConversionError::InvalidUuid), + Err(e) => return Err(RecordConversionError::OperationConversionError(e)), }; - let (embedding, encoding) = match proto_submit.vector { + let (embedding, encoding) = match operation_record_proto.vector { Some(proto_vector) => match proto_vector.try_into() { Ok((embedding, encoding)) => (Some(embedding), Some(encoding)), - Err(e) => return Err(EmbeddingRecordConversionError::VectorConversionError(e)), + Err(e) => return Err(RecordConversionError::VectorConversionError(e)), }, // If there is no vector, there is no encoding None => (None, None), }; - let metadata: Option = match proto_submit.metadata { + let metadata: Option = match operation_record_proto.metadata { Some(proto_metadata) => match proto_metadata.try_into() { Ok(metadata) => Some(metadata), - Err(e) => { - return Err( - EmbeddingRecordConversionError::UpdateMetadataValueConversionError(e), - ) - } + Err(e) => return Err(RecordConversionError::UpdateMetadataValueConversionError(e)), }, None => None, }; - Ok(EmbeddingRecord { - id: proto_submit.id, - seq_id: seq_id, - embedding: embedding, - encoding: encoding, - metadata: metadata, - operation: op, - collection_id: collection_uuid, + Ok(OperationRecord { + id: operation_record_proto.id, + embedding, + encoding, + metadata, + operation, + }) + } +} + +impl TryFrom for LogRecord { + type Error = RecordConversionError; + + fn try_from(log_record_proto: chroma_proto::LogRecord) -> Result { + let record = match log_record_proto.record { + Some(proto_record) => OperationRecord::try_from(proto_record)?, + None => { + return Err(RecordConversionError::DecodeError( + ConversionError::DecodeError, + )) + } + }; + Ok(LogRecord { + log_offset: log_record_proto.log_offset, + record, }) } } @@ -200,7 +207,6 @@ Vector Embedding Record #[derive(Debug)] pub(crate) struct VectorEmbeddingRecord { pub(crate) id: String, - pub(crate) seq_id: SeqId, pub(crate) vector: Vec, } @@ -213,19 +219,16 @@ Vector Query Result #[derive(Debug)] pub(crate) struct VectorQueryResult { pub(crate) id: String, - pub(crate) seq_id: SeqId, pub(crate) distance: f32, pub(crate) vector: Option>, } #[cfg(test)] mod tests { - use std::collections::HashMap; - - use num_bigint::BigInt; - use super::*; use crate::{chroma_proto, types::UpdateMetadataValue}; + use std::collections::HashMap; + use uuid::Uuid; fn as_byte_view(input: &[f32]) -> Vec { unsafe { @@ -238,7 +241,7 @@ mod tests { } #[test] - fn test_embedding_record_try_from() { + fn test_operation_record_try_from() { let mut metadata = chroma_proto::UpdateMetadata { metadata: HashMap::new(), }; @@ -253,29 +256,68 @@ mod tests { encoding: chroma_proto::ScalarEncoding::Float32 as i32, dimension: 3, }; - let proto_submit = chroma_proto::SubmitEmbeddingRecord { + let proto_submit = chroma_proto::OperationRecord { id: "00000000-0000-0000-0000-000000000000".to_string(), vector: Some(proto_vector), metadata: Some(metadata), operation: chroma_proto::Operation::Add as i32, - collection_id: "00000000-0000-0000-0000-000000000000".to_string(), }; - let converted_embedding_record: EmbeddingRecord = - EmbeddingRecord::try_from((proto_submit, BigInt::from(42))).unwrap(); - assert_eq!(converted_embedding_record.id, Uuid::nil().to_string()); - assert_eq!(converted_embedding_record.seq_id, BigInt::from(42)); + let converted_operation_record = OperationRecord::try_from(proto_submit).unwrap(); + assert_eq!(converted_operation_record.id, Uuid::nil().to_string()); + assert_eq!( + converted_operation_record.embedding, + Some(vec![1.0, 2.0, 3.0]) + ); + assert_eq!( + converted_operation_record.encoding, + Some(ScalarEncoding::FLOAT32) + ); + let metadata = converted_operation_record.metadata.unwrap(); + assert_eq!(metadata.len(), 1); + assert_eq!(metadata.get("foo").unwrap(), &UpdateMetadataValue::Int(42)); + assert_eq!(converted_operation_record.operation, Operation::Add); + } + + #[test] + fn test_log_record_try_from_record_log() { + let mut metadata = chroma_proto::UpdateMetadata { + metadata: HashMap::new(), + }; + metadata.metadata.insert( + "foo".to_string(), + chroma_proto::UpdateMetadataValue { + value: Some(chroma_proto::update_metadata_value::Value::IntValue(42)), + }, + ); + let proto_vector = chroma_proto::Vector { + vector: as_byte_view(&[1.0, 2.0, 3.0]), + encoding: chroma_proto::ScalarEncoding::Float32 as i32, + dimension: 3, + }; + let proto_submit = chroma_proto::OperationRecord { + id: "00000000-0000-0000-0000-000000000000".to_string(), + vector: Some(proto_vector), + metadata: Some(metadata), + operation: chroma_proto::Operation::Add as i32, + }; + let record_log = chroma_proto::LogRecord { + log_offset: 42, + record: Some(proto_submit), + }; + let converted_log_record = LogRecord::try_from(record_log).unwrap(); + assert_eq!(converted_log_record.record.id, Uuid::nil().to_string()); + assert_eq!(converted_log_record.log_offset, 42); assert_eq!( - converted_embedding_record.embedding, + converted_log_record.record.embedding, Some(vec![1.0, 2.0, 3.0]) ); assert_eq!( - converted_embedding_record.encoding, + converted_log_record.record.encoding, Some(ScalarEncoding::FLOAT32) ); - let metadata = converted_embedding_record.metadata.unwrap(); + let metadata = converted_log_record.record.metadata.unwrap(); assert_eq!(metadata.len(), 1); assert_eq!(metadata.get("foo").unwrap(), &UpdateMetadataValue::Int(42)); - assert_eq!(converted_embedding_record.operation, Operation::Add); - assert_eq!(converted_embedding_record.collection_id, Uuid::nil()); + assert_eq!(converted_log_record.record.operation, Operation::Add); } } diff --git a/rust/worker/src/types/scalar_encoding.rs b/rust/worker/src/types/scalar_encoding.rs index afcaf6b2e30..5114328886a 100644 --- a/rust/worker/src/types/scalar_encoding.rs +++ b/rust/worker/src/types/scalar_encoding.rs @@ -5,7 +5,7 @@ use crate::{ }; use thiserror::Error; -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) enum ScalarEncoding { FLOAT32, INT32, diff --git a/rust/worker/src/types/segment.rs b/rust/worker/src/types/segment.rs index 02e3dcd4434..6a0eb7aa76a 100644 --- a/rust/worker/src/types/segment.rs +++ b/rust/worker/src/types/segment.rs @@ -3,22 +3,24 @@ use crate::{ chroma_proto, errors::{ChromaError, ErrorCodes}, }; +use std::collections::HashMap; +use std::vec::Vec; use thiserror::Error; use uuid::Uuid; -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) enum SegmentType { HnswDistributed, } -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) struct Segment { pub(crate) id: Uuid, pub(crate) r#type: SegmentType, pub(crate) scope: SegmentScope, - pub(crate) topic: Option, pub(crate) collection: Option, pub(crate) metadata: Option, + pub(crate) file_path: HashMap>, } #[derive(Error, Debug)] @@ -48,6 +50,8 @@ impl TryFrom for Segment { type Error = SegmentConversionError; fn try_from(proto_segment: chroma_proto::Segment) -> Result { + let mut proto_segment = proto_segment; + let segment_uuid = match Uuid::try_parse(&proto_segment.id) { Ok(uuid) => uuid, Err(_) => return Err(SegmentConversionError::InvalidUuid), @@ -79,13 +83,19 @@ impl TryFrom for Segment { } }; + let mut file_paths = HashMap::new(); + let drain = proto_segment.file_paths.drain(); + for (key, value) in drain { + file_paths.insert(key, value.paths); + } + Ok(Segment { id: segment_uuid, r#type: segment_type, scope: scope, - topic: proto_segment.topic, collection: collection_uuid, metadata: segment_metadata, + file_path: file_paths, }) } } @@ -112,15 +122,14 @@ mod tests { id: "00000000-0000-0000-0000-000000000000".to_string(), r#type: "urn:chroma:segment/vector/hnsw-distributed".to_string(), scope: chroma_proto::SegmentScope::Vector as i32, - topic: Some("test".to_string()), collection: Some("00000000-0000-0000-0000-000000000000".to_string()), metadata: Some(metadata), + file_paths: HashMap::new(), }; let converted_segment: Segment = proto_segment.try_into().unwrap(); assert_eq!(converted_segment.id, Uuid::nil()); assert_eq!(converted_segment.r#type, SegmentType::HnswDistributed); assert_eq!(converted_segment.scope, SegmentScope::VECTOR); - assert_eq!(converted_segment.topic, Some("test".to_string())); assert_eq!(converted_segment.collection, Some(Uuid::nil())); let metadata = converted_segment.metadata.unwrap(); assert_eq!(metadata.len(), 1); diff --git a/rust/worker/src/types/segment_scope.rs b/rust/worker/src/types/segment_scope.rs index d2c1fb5392f..1777f480b41 100644 --- a/rust/worker/src/types/segment_scope.rs +++ b/rust/worker/src/types/segment_scope.rs @@ -5,7 +5,7 @@ use crate::{ }; use thiserror::Error; -#[derive(Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub(crate) enum SegmentScope { VECTOR, METADATA, diff --git a/rust/worker/src/types/types.rs b/rust/worker/src/types/types.rs index e87337cc511..ee6368c6282 100644 --- a/rust/worker/src/types/types.rs +++ b/rust/worker/src/types/types.rs @@ -1,5 +1,4 @@ use crate::errors::{ChromaError, ErrorCodes}; -use num_bigint::BigInt; use thiserror::Error; /// A macro for easily implementing match arms for a base error type with common errors. @@ -32,5 +31,3 @@ impl ChromaError for ConversionError { } } } - -pub(crate) type SeqId = BigInt; diff --git a/server.htpasswd b/server.htpasswd deleted file mode 100644 index 77f277a399b..00000000000 --- a/server.htpasswd +++ /dev/null @@ -1 +0,0 @@ -admin:$2y$05$e5sRb6NCcSH3YfbIxe1AGu2h5K7OOd982OXKmd8WyQ3DRQ4MvpnZS